From d81a5719ada63b602e2020c82a991a6a4c855185 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:56:05 -0700 Subject: [PATCH] fix(cache): keep cache hits free of unset fields, and follow intentional behavior changes in CircleCI suites (#46359) * test(integration): follow count_tokens error surfacing and decisions pricing in CircleCI suites Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): restore Responses input_file PDF pricing coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): keep integration AGENTS guide unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cache): preserve provider field presence on cache hits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect credential display name in response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): align OpenInference refusal field expectations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect Akto request ingestion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): expect missing tool error wording Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect no refusal field in cursor chat wire Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching.py | 2 +- litellm/caching/caching_handler.py | 2 +- .../authorization/test_team_admin_gate.py | 1 + .../mcp/test_mcp_accounting_guardrails.py | 4 +- .../test_pdf_document_count_tokens_chaos.py | 98 ++++++------ .../test_pdf_document_count_tokens_wire.py | 143 ++++++++++-------- .../observability/_openinference_support.py | 1 - .../test_apply_guardrail_endpoint.py | 4 +- .../providers/test_cursor_responses_wire.py | 1 - .../test_decisions_openai_format_wire.py | 12 +- .../integration/spend/test_cache_and_quota.py | 63 +++++++- tests/unit/caching/test_caching_handler.py | 47 ++++++ 12 files changed, 255 insertions(+), 123 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 05275d58c7d..1d101b65e36 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -709,7 +709,7 @@ class Cache: cache_key = self.get_cache_key(**kwargs) if cache_key is not None: if isinstance(result, BaseModel): - result = result.model_dump_json() + result = result.model_dump_json(exclude_unset=True) ## DEFAULT TTL ## if self.ttl is not None: diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index c97baae293a..3fb0d0308d9 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -1180,7 +1180,7 @@ class LLMCachingHandler: ) ) else: - result_json: Final = result.model_dump_json() + result_json: Final = result.model_dump_json(exclude_unset=True) create_cache_write_task( lambda: cache.async_add_cache( result_json, diff --git a/tests/integration/authorization/test_team_admin_gate.py b/tests/integration/authorization/test_team_admin_gate.py index 7b4fa52c4e4..a7513277609 100644 --- a/tests/integration/authorization/test_team_admin_gate.py +++ b/tests/integration/authorization/test_team_admin_gate.py @@ -811,6 +811,7 @@ def test_only_proxy_admin_reads_or_changes_credentials( assert credential_for_model.status_code == 200, credential_for_model.text assert credential_for_model.json() == { "credential_name": f"{model_name}-credential-{model_id}", + "display_name": None, "credential_values": {"api_base": f"{s.gateway.upstream_url}/v1"}, "credential_info": {}, }, credential_for_model.text diff --git a/tests/integration/mcp/test_mcp_accounting_guardrails.py b/tests/integration/mcp/test_mcp_accounting_guardrails.py index 0d1c707d1ff..72d6c860d2f 100644 --- a/tests/integration/mcp/test_mcp_accounting_guardrails.py +++ b/tests/integration/mcp/test_mcp_accounting_guardrails.py @@ -628,7 +628,7 @@ def test_a_schema_past_the_scan_depth_is_not_published_while_one_at_the_limit_is assert scanned["texts"] == ["deepest leaf", "r-1"], scanned unpublished: Final = caller.call(f"{alias}-deep", {"record": "r-2"}) assert _worker(pinned) == worker - assert unpublished.error is not None and "Tool 'deep' not found" in unpublished.raw, unpublished.raw + assert unpublished.error == "Error: Tool deep not found", unpublished.raw reached: Final = tool_calls(peer.drain()) assert [object_value(JSON_OBJECT.validate_python(call["body"])["params"])["name"] for call in reached] == [ "shallow" @@ -643,7 +643,7 @@ def test_a_schema_past_the_scan_depth_is_not_published_while_one_at_the_limit_is ], rows assert [_tool_metadata(row)["name"] for row in rows[1:]] == ["shallow", "deep"], rows failed: Final = object_value(JSON_OBJECT.validate_python(rows[2]["metadata"])["error_information"]) - assert failed["error_message"] == "404: Tool 'deep' not found", failed + assert failed["error_message"] == "Tool deep not found", failed def test_an_adapter_returning_the_wrong_number_of_texts_fails_closed_before_the_peer(hooks_rig: HooksRig) -> None: diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py index 80ace4c4535..63325e7a005 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py @@ -15,8 +15,6 @@ import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.pdf_document import ( - COUNT_REFUSED, - COUNT_TOKENS_TARGET, LETTER, messages_body, pdf_document, @@ -42,7 +40,7 @@ class _Counted: pages: int status: int text: str - input_tokens: int | None + total_tokens: int | None @dataclass(frozen=True, slots=True) @@ -54,8 +52,6 @@ class _Sent: def _peer(request: Request) -> Reply: - if request.target == COUNT_TOKENS_TARGET: - return COUNT_REFUSED return anthropic_peer(request) @@ -79,31 +75,31 @@ def _count_body(pages: int) -> dict[str, JsonValue]: return {"model": _MODEL, "messages": messages_body(_MODEL, blocks, _ASK)["messages"]} -def _input_tokens(response: httpx.Response) -> int | None: +def _total_tokens(response: httpx.Response) -> int | None: if response.status_code != 200: return None - counted: Final = _JSON_OBJECT.validate_json(response.content).get("input_tokens") + counted: Final = _JSON_OBJECT.validate_json(response.content).get("total_tokens") return counted if isinstance(counted, int) else None -async def _fire_counts(url: str, key: str, *, tolerate_transport_errors: bool = False) -> tuple[_Counted, ...]: +async def _fire_counts(url: str, key: str) -> tuple[_Counted, ...]: async def one(client: httpx.AsyncClient, index: int) -> _Counted: pages: Final = _PAGE_COUNTS[index % len(_PAGE_COUNTS)] response: Final = await client.post( - "/v1/messages/count_tokens", json=_count_body(pages), headers={"Authorization": f"Bearer {key}"} + "/utils/token_counter", + params={"call_endpoint": "false"}, + json=_count_body(pages), + headers={"Authorization": f"Bearer {key}"}, ) - return _Counted(pages, response.status_code, response.text, _input_tokens(response)) + return _Counted(pages, response.status_code, response.text, _total_tokens(response)) async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - results: Final = await asyncio.gather( - *(one(client, index) for index in range(_COUNT_BURST)), return_exceptions=tolerate_transport_errors - ) - for result in results: - assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) - return tuple(result for result in results if isinstance(result, _Counted)) + return tuple(await asyncio.gather(*(one(client, index) for index in range(_COUNT_BURST)))) -async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: +async def _fire_messages( + url: str, key: str, *, tolerate_transport_errors: bool = False +) -> tuple[_Sent, ...]: async def one(client: httpx.AsyncClient) -> _Sent: marker: Final = uuid.uuid4().hex body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") @@ -111,13 +107,20 @@ async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: return _Sent(marker, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - return tuple(await asyncio.gather(*(one(client) for _ in range(_MESSAGE_BURST)))) + results: Final = await asyncio.gather( + *(one(client) for _ in range(_MESSAGE_BURST)), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Sent)) def _baseline(gateway: Gateway) -> int: - response: Final = gateway.request("POST", "/v1/messages/count_tokens", _count_body(0)) + response: Final = gateway.request( + "POST", "/utils/token_counter", _count_body(0), params={"call_endpoint": "false"} + ) assert response.status_code == 200, response.text - counted: Final = _input_tokens(response) + counted: Final = _total_tokens(response) assert counted is not None, response.text return counted @@ -125,11 +128,11 @@ def _baseline(gateway: Gateway) -> int: def _assert_exact(counts: tuple[_Counted, ...], baseline: int) -> None: for item in counts: assert item.status == 200, (item.pages, item.status, item.text) - assert item.input_tokens is not None and item.input_tokens - baseline == rendered_tokens( + assert item.total_tokens is not None and item.total_tokens - baseline == rendered_tokens( (LETTER,) * item.pages ), ( item.pages, - item.input_tokens, + item.total_tokens, baseline, ) @@ -155,7 +158,7 @@ def _held_upstream_connections(pid: int, upstream: str) -> int: @pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( +async def test_pdf_local_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( gateway: Gateway, tmp_path: Path ) -> None: with wire_server(_peer) as wire: @@ -174,8 +177,8 @@ async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_eac ) received: Final = wire.drain() targets: Final = [request.target for request in received] - assert targets.count(COUNT_TOKENS_TARGET) == _COUNT_BURST, targets assert targets.count("/v1/messages") == _MESSAGE_BURST, targets + assert len(targets) == _MESSAGE_BURST, targets _assert_exact(counts, baseline) assert {item.marker for item in messages} == { _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in messages if item.status == 200 @@ -185,7 +188,7 @@ async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_eac @pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pages( +async def test_worker_sigkill_mid_pdf_message_burst_preserves_local_page_pricing_and_single_logs( gateway: Gateway, tmp_path: Path ) -> None: release: Final = threading.Event() @@ -198,29 +201,38 @@ async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pag seconds=graceful_stop_seconds(), ) owned_url: Final = str(owned.gateway.client.base_url) - burst: Final = asyncio.create_task( - _fire_counts(owned_url, owned.gateway.key, tolerate_transport_errors=True) - ) - try: - await asyncio.to_thread( - eventually, lambda: wire.received.qsize(), lambda size: size >= _COUNT_BURST, 90 + async with asyncio.TaskGroup() as tasks: + burst = tasks.create_task( + _fire_messages(owned_url, owned.gateway.key, tolerate_transport_errors=True) ) - held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) - victim: Final = max(workers, key=held_by.__getitem__) - with owned.replacing_worker(victim) as kill: - await asyncio.to_thread(kill) + try: + await asyncio.to_thread( + eventually, lambda: wire.received.qsize(), lambda size: size >= _MESSAGE_BURST, 90 + ) + held_by: Final = MappingProxyType( + {pid: _held_upstream_connections(pid, wire.url) for pid in workers} + ) + victim: Final = max(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + message_requests: Final = wire.drain() + finally: release.set() - served: Final = await burst - wire.drain() - finally: - release.set() baseline: Final = _baseline(owned.gateway) after: Final = await _fire_counts(owned_url, owned.gateway.key) after_received: Final = wire.drain() - assert sum(held_by.values()) == _COUNT_BURST, held_by + assert sum(held_by.values()) == _MESSAGE_BURST, held_by assert held_by[victim] > 0, held_by - assert len(served) == _COUNT_BURST - held_by[victim], (len(served), held_by) + assert len(message_requests) == _MESSAGE_BURST, message_requests + assert all(request.target == "/v1/messages" for request in message_requests), message_requests + assert len(served) == _MESSAGE_BURST - held_by[victim], (len(served), held_by) assert len(after) == _COUNT_BURST, len(after) - assert [request.target for request in after_received] == [COUNT_TOKENS_TARGET] * (_COUNT_BURST + 1) - _assert_exact(served, baseline) + assert after_received == (), after_received + assert {item.marker for item in served} == { + _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in served if item.status == 200 + } + for item in served: + _single_spend_row(item) _assert_exact(after, baseline) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py index 41d5f3d8685..8eeaab66553 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py @@ -22,8 +22,8 @@ from integration._support.pdf_document import ( pdf_bytes, pdf_data_url, pdf_document, - rendered_tokens, responses_input_file, + rendered_tokens, text_document, ) from integration._support.provider import SharedProvider @@ -70,31 +70,33 @@ def _turn(blocks: Sequence[JsonValue]) -> list[JsonValue]: return messages_body(_MODEL, blocks, _ASK)["messages"] -def _count(gateway: Gateway, provider: SharedProvider, blocks: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) - response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": _MODEL, "messages": _turn(blocks)}) - assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] - return _int(_JSON_OBJECT.validate_json(response.content)["input_tokens"]) - - -def _local(gateway: Gateway, parts: Sequence[JsonValue]) -> int: +def _local(gateway: Gateway, provider: SharedProvider, parts: Sequence[JsonValue]) -> int: response: Final = gateway.request( "POST", "/utils/token_counter", {"model": _MODEL, "messages": _turn(parts)}, params={"call_endpoint": "false"} ) assert response.status_code == 200, response.text + assert provider.received() == () return _int(_JSON_OBJECT.validate_json(response.content)["total_tokens"]) -def _responses_count(gateway: Gateway, provider: SharedProvider, items: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) +def _responses_count(gateway: Gateway, model: str, pages: Sequence[Page]) -> int: response: Final = gateway.request( "POST", "/v1/responses/input_tokens", - {"model": _MODEL, "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": _ASK}]}]}, + { + "model": model, + "input": [ + { + "role": "user", + "content": [ + responses_input_file(pdf_data_url(pages)), + {"type": "input_text", "text": _ASK}, + ], + } + ], + }, ) assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["object"] == "response.input_tokens", response.text return _int(payload["input_tokens"]) @@ -144,108 +146,123 @@ def _prompt_tokens(call_id: str) -> int: @pytest.mark.parametrize("pages", [1, 3, 12]) -def test_count_tokens_fallback_prices_each_blank_page_as_a_rendered_image( +def test_utils_token_counter_prices_each_blank_page_as_a_rendered_image( gateway: Gateway, provider: SharedProvider, pages: int ) -> None: letter_pages: Final = (LETTER,) * pages - with_document: Final = _count(gateway, provider, [pdf_document(letter_pages)]) - without: Final = _count(gateway, provider, []) + with_document: Final = _local(gateway, provider, [pdf_document(letter_pages)]) + without: Final = _local(gateway, provider, []) assert with_document - without == rendered_tokens(letter_pages), (with_document, without) -def test_count_tokens_fallback_adds_a_page_text_to_its_rendered_image( +def test_utils_token_counter_adds_a_page_text_to_its_rendered_image( gateway: Gateway, provider: SharedProvider ) -> None: - text_page: Final = _count(gateway, provider, [pdf_document((Page(text=_TEXT),))]) - blank_page: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - text_source: Final = _count(gateway, provider, [text_document(_TEXT)]) - without: Final = _count(gateway, provider, []) + text_page: Final = _local(gateway, provider, [pdf_document((Page(text=_TEXT),))]) + blank_page: Final = _local(gateway, provider, [pdf_document((LETTER,))]) + text_source: Final = _local(gateway, provider, [text_document(_TEXT)]) + without: Final = _local(gateway, provider, []) assert text_source - without > 0, (text_source, without) assert text_page - blank_page == text_source - without, (text_page, blank_page, text_source, without) -def test_count_tokens_fallback_prices_a_page_by_its_rendered_area(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_prices_a_page_by_its_rendered_area( + gateway: Gateway, provider: SharedProvider +) -> None: mixed: Final = (LETTER, NARROW, LETTER) - without: Final = _count(gateway, provider, []) - narrow: Final = _count(gateway, provider, [pdf_document((NARROW,))]) + without: Final = _local(gateway, provider, []) + narrow: Final = _local(gateway, provider, [pdf_document((NARROW,))]) assert narrow - without == rendered_tokens((NARROW,)), (narrow, without) - assert _count(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) + assert _local(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) -def test_count_tokens_fallback_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: - both: Final = _count(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) - without: Final = _count(gateway, provider, []) +def test_utils_token_counter_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: + both: Final = _local(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) + without: Final = _local(gateway, provider, []) assert both - without == rendered_tokens((LETTER,) * 4), (both, without) -def test_count_tokens_fallback_prices_a_pdf_without_pages_as_nothing( - gateway: Gateway, provider: SharedProvider -) -> None: - assert _count(gateway, provider, [pdf_document(())]) == _count(gateway, provider, []) +def test_utils_token_counter_prices_a_pdf_without_pages_as_nothing(gateway: Gateway, provider: SharedProvider) -> None: + assert _local(gateway, provider, [pdf_document(())]) == _local(gateway, provider, []) -def test_count_tokens_fallback_adds_title_and_context_once_per_document( +def test_utils_token_counter_adds_title_and_context_once_per_document( gateway: Gateway, provider: SharedProvider ) -> None: - without: Final = _count(gateway, provider, []) - plain: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - annotated: Final = _count(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) - title: Final = _count(gateway, provider, [text_document(_TITLE)]) - without - context: Final = _count(gateway, provider, [text_document(_CONTEXT)]) - without + without: Final = _local(gateway, provider, []) + plain: Final = _local(gateway, provider, [pdf_document((LETTER,))]) + annotated: Final = _local(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) + title: Final = _local(gateway, provider, [text_document(_TITLE)]) - without + context: Final = _local(gateway, provider, [text_document(_CONTEXT)]) - without assert title > 0 and context > 0, (title, context) assert annotated - plain == title + context, (annotated, plain, title, context) @pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_count_tokens_fallback_prices_unreadable_pdf_bytes_like_one_image( +def test_utils_token_counter_prices_unreadable_pdf_bytes_like_one_image( gateway: Gateway, provider: SharedProvider, label: str ) -> None: data: Final = _UNREADABLE[label] - as_pdf: Final = _count(gateway, provider, [document(base64_source(data))]) - as_png: Final = _count(gateway, provider, [document(base64_source(data, media_type="image/png"))]) - without: Final = _count(gateway, provider, []) + as_pdf: Final = _local(gateway, provider, [document(base64_source(data))]) + as_png: Final = _local(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + without: Final = _local(gateway, provider, []) assert as_pdf == as_png, (as_pdf, as_png) assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) -def test_count_tokens_fallback_reads_a_list_wrapped_base64_string_like_the_bare_string( +def test_utils_token_counter_reads_a_list_wrapped_base64_string_like_the_bare_string( gateway: Gateway, provider: SharedProvider ) -> None: raw: Final = encoded(pdf_bytes((LETTER,))) - wrapped: Final = _count(gateway, provider, [document(base64_source([raw]))]) - assert wrapped == _count(gateway, provider, [document(base64_source(raw))]), wrapped + wrapped: Final = _local(gateway, provider, [document(base64_source([raw]))]) + assert wrapped == _local(gateway, provider, [document(base64_source(raw))]), wrapped -def test_count_tokens_fallback_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: pages: Final = (LETTER, LETTER) locked: Final = document(base64_source(encoded(_locked(pdf_bytes(pages), "")))) - assert _count(gateway, provider, [locked]) - _count(gateway, provider, []) == rendered_tokens(pages) + assert _local(gateway, provider, [locked]) - _local(gateway, provider, []) == rendered_tokens(pages) -def test_count_tokens_fallback_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: raw: Final = encoded(pdf_bytes((LETTER,))) - without: Final = _count(gateway, provider, []) - labelled: Final = _count(gateway, provider, [document(base64_source(raw))]) + without: Final = _local(gateway, provider, []) + labelled: Final = _local(gateway, provider, [document(base64_source(raw))]) assert labelled - without == rendered_tokens((LETTER,)), (labelled, without) - mislabelled: Final = _count(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) - assert mislabelled == _count(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) + mislabelled: Final = _local(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) + assert mislabelled == _local(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) -def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: - one: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,)))]) - three: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,) * 3))]) +def test_messages_count_tokens_surfaces_provider_refusal_without_returning_tokens( + gateway: Gateway, provider: SharedProvider +) -> None: + provider.expect(COUNT_REFUSED) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": _MODEL, "messages": _turn([pdf_document((LETTER,))])}, + ) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text + assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] + + +def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway, provider: SharedProvider) -> None: + one: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,)))]) + three: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,) * 3))]) assert three - one == rendered_tokens((LETTER,) * 2), (three, one) - text_page: Final = _local(gateway, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) - text_part: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) + text_page: Final = _local(gateway, provider, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) + text_part: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) assert text_part - one > 0, (text_part, one) assert text_page - one == text_part - one, (text_page, text_part, one) -def test_responses_input_tokens_fallback_prices_an_input_file_by_its_pages( - gateway: Gateway, provider: SharedProvider -) -> None: - one: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,)))]) - three: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,) * 3))]) +def test_responses_input_tokens_prices_an_input_file_by_its_pages(gateway: Gateway, provider: SharedProvider) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model="xai/grok-4.20", api_base=f"{provider.wire.url}/v1") + one: Final = _responses_count(gateway, model, (LETTER,)) + three: Final = _responses_count(gateway, model, (LETTER,) * 3) + assert provider.received() == () assert three - one == rendered_tokens((LETTER,) * 2), (three, one) diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index ee6dee07c46..61def2a355b 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -204,7 +204,6 @@ def _chat_caller_message(message: dict[str, JsonValue]) -> dict[str, JsonValue]: return { **{key: value for key, value in message.items() if key != "tool_calls"}, **({"tool_calls": [_chat_caller_tool_call(call) for call in tool_calls]} if tool_calls else {}), - "provider_specific_fields": {"refusal": None}, } diff --git a/tests/integration/observability/test_apply_guardrail_endpoint.py b/tests/integration/observability/test_apply_guardrail_endpoint.py index 4ff11a50667..9fa7d1a74db 100644 --- a/tests/integration/observability/test_apply_guardrail_endpoint.py +++ b/tests/integration/observability/test_apply_guardrail_endpoint.py @@ -212,8 +212,8 @@ def test_apply_guardrail_forwards_messages_to_the_guardrail(gateway: Gateway) -> requests: Final = akto.drain() assert [(request.method, request.target) for request in requests] == [ - ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true"), - ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true"), + ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true&ingest_data=true"), + ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true&ingest_data=true"), ] assert [request.headers["authorization"] for request in requests] == ["synthetic-akto-apply-key"] * 2 bodies: Final = tuple( diff --git a/tests/integration/providers/test_cursor_responses_wire.py b/tests/integration/providers/test_cursor_responses_wire.py index cf766c9105a..6e8c877f7ca 100644 --- a/tests/integration/providers/test_cursor_responses_wire.py +++ b/tests/integration/providers/test_cursor_responses_wire.py @@ -399,7 +399,6 @@ def test_cursor_messages_body_uses_chat_completions_wire( "message": { "content": "Patch inspected.", "role": "assistant", - "provider_specific_fields": {"refusal": None}, }, "provider_specific_fields": {}, } diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py index d989a0c49aa..cb021e02e88 100644 --- a/tests/integration/providers/test_decisions_openai_format_wire.py +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -285,9 +285,17 @@ def _expected_spend(provider: _Provider) -> float: return 0.0 cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) prices: Final = object_value(cost_map[provider.cost_map_key]) - return _INPUT_TOKENS * _number(prices["input_cost_per_token"]) + _OUTPUT_TOKENS * _number( - prices["output_cost_per_token"] + input_key: Final = ( + "input_cost_per_token_decisions" + if "input_cost_per_token_decisions" in prices + else "input_cost_per_token" ) + output_key: Final = ( + "output_cost_per_token_decisions" + if "output_cost_per_token_decisions" in prices + else "output_cost_per_token" + ) + return _INPUT_TOKENS * _number(prices[input_key]) + _OUTPUT_TOKENS * _number(prices[output_key]) def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle: diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1cc03838f8d..d02ebccf710 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -16,7 +16,9 @@ from integration._support.database import read_rows, scratch_database from integration._support.database_relay import database_relay from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests from integration._support.process import owned_proxy +from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request, wire_server +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse @pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting") @@ -307,7 +309,43 @@ def test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_ gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, ): - model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + counting_upstream: Final = register_scenario( + f"budget-count-{uuid.uuid4().hex}", + RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /v1/responses/input_tokens": JsonResponse( + content_type="application/json", body={"input_tokens": 12} + ), + "POST /v1/chat/completions": JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1729380000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! This is a mock response from the fake OpenAI endpoint.", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, + }, + ), + }, + ), + ) + scenario.cleanups.callback(delete_scenario, counting_upstream) + model: Final = scenario.model( + api_base=f"{counting_upstream.api_base()}/v1", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) key: Final = scenario.key(models=[model], max_budget=0.1) digest: Final = sha256(key.encode()).hexdigest() upstream.get("/__observations").raise_for_status() @@ -325,19 +363,30 @@ def test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_ assert count.status_code == 200, count.text assert count.json() == counts[0].json(), count.text input_tokens: Final = counts[0].json()["input_tokens"] - assert isinstance(input_tokens, int) and input_tokens > 0, counts[0].text - assert upstream.get("/__observations").json()["requests"] == [] + assert input_tokens == 12, counts[0].text + counted_requests: Final = upstream.get("/__observations").json()["requests"] + assert counted_requests == [ + { + "path": f"/{counting_upstream.scenario_id}/v1/responses/input_tokens", + "authorization": "Bearer integration-provider-key", + "body": {"model": "gpt-4o-mini", "input": [{"role": "user", "content": "hello!!!"}]}, + "method": "POST", + "api_key": "", + } + ] * 3, counted_requests + completion_prompt: Final = f"after counting {uuid.uuid4().hex}" completion: Final = gateway.request( "POST", "/v1/chat/completions", - {"model": model, "messages": [{"role": "user", "content": f"after counting {uuid.uuid4().hex}"}]}, + {"model": model, "messages": [{"role": "user", "content": completion_prompt}]}, key=key, ) assert completion.status_code == 200, completion.text assert completion.json()["usage"]["total_tokens"] == 40, completion.text - assert [request["path"] for request in upstream.get("/__observations").json()["requests"]] == [ - "/v1/chat/completions" - ] + completion_requests: Final = upstream.get("/__observations").json()["requests"] + assert len(completion_requests) == 1, completion_requests + assert completion_requests[0]["path"] == f"/{counting_upstream.scenario_id}/v1/chat/completions" + assert completion_requests[0]["body"]["messages"] == [{"role": "user", "content": completion_prompt}] spent: Final = eventually( lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) > 0, diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 2520f2eb2da..85e677aa990 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -46,6 +46,7 @@ from typing import Final, cast from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm._logging import verbose_logger +from litellm.proxy.utils import model_dump_with_preserved_fields import logging import json import httpx @@ -2706,6 +2707,52 @@ async def test_async_get_cache_heals_stored_completion_without_choices(): assert healed.choices[0].message.content == chat_completion_response.choices[0].message.content +@pytest.mark.parametrize("explicit_nulls", [False, True]) +@pytest.mark.asyncio +async def test_async_cache_preserves_presence_of_provider_optional_fields(explicit_nulls: bool) -> None: + setup_cache() + handler: Final = LLMCachingHandler(original_function=litellm.acompletion, request_kwargs={}, start_time=_FIXED_START) + kwargs: Final = {"model": "gpt-4o", "messages": _unique_messages(), "caching": True} + message: Final = {"role": "assistant", "content": "cache response"} | ( + {"refusal": None} if explicit_nulls else {} + ) + choice: Final = {"index": 0, "finish_reason": "stop", "message": message} | ( + {"logprobs": None} if explicit_nulls else {} + ) + response_payload: Final = { + "id": "cache-field-round-trip", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [choice], + **({"system_fingerprint": None} if explicit_nulls else {}), + } + response: Final = ModelResponse.model_validate(response_payload) + await handler.async_set_cache(result=response, original_function=litellm.acompletion, kwargs=kwargs) + await asyncio.gather(*_PENDING_CACHE_WRITES) + lookup: Final = await handler.async_get_cache( + model="gpt-4o", + original_function=litellm.acompletion, + logging_obj=_completion_logging_obj(CallTypes.acompletion.value), + start_time=_FIXED_START, + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + ) + cached_response: Final = lookup.cached_result + assert isinstance(cached_response, ModelResponse) + original_body: Final = model_dump_with_preserved_fields(response) + cached_body: Final = model_dump_with_preserved_fields(cached_response) + assert ("system_fingerprint" in original_body) is explicit_nulls + assert ("system_fingerprint" in cached_body) is explicit_nulls + original_choices: Final = cast(list[dict[str, object]], original_body["choices"]) + cached_choices: Final = cast(list[dict[str, object]], cached_body["choices"]) + original_message: Final = cast(dict[str, object], original_choices[0]["message"]) + cached_message: Final = cast(dict[str, object], cached_choices[0]["message"]) + assert ("refusal" in original_message) is explicit_nulls + assert ("refusal" in cached_message) is explicit_nulls + assert original_choices[0].get("logprobs") == cached_choices[0].get("logprobs") + + def test_sync_set_cache_skips_response_without_output(): setup_cache() handler = LLMCachingHandler(original_function=completion, request_kwargs={}, start_time=_FIXED_START)