mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(count_tokens): keep the Mantle config for a Converse Bedrock deployment's count
A bedrock/converse/ deployment whose count goes to Mantle built its count options with the Bedrock invoke config, because every bedrock deployment got that config there. Converse has no messages transform, so the count fell back to the invoke config on the unprefixed count model. It inlined an output_config.format schema into the user message and sent betas the bedrock row allows but Converse inference drops, so the count came out lower than what inference bills A Mantle count now uses the invoke config only for a bedrock deployment whose /v1/messages runs on invoke. Converse deployments, bedrock_mantle deployments and counts with no deployment model keep the Mantle config, which sends output_config natively the way Converse inference does The count resolves the deployment route once and reuses it for the chat translator and the options config, and it reads litellm_params as a typed mapping instead of an Any Tests now cover a Converse deployment's Mantle count against its own Converse inference, and Mantle counts for a bedrock_mantle deployment, one named through custom_llm_provider, and one with no deployment. The proxy retry rows check which host got each count again
This commit is contained in:
parent
77f0f3104a
commit
64c658486c
4 changed files with 173 additions and 52 deletions
|
|
@ -12,6 +12,7 @@ from typing_extensions import assert_never
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
|
|
@ -79,8 +80,11 @@ def _deployment_route(litellm_params: Mapping[str, object]) -> tuple[str, str] |
|
|||
return model, provider
|
||||
|
||||
|
||||
def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig:
|
||||
route: Final = _deployment_route(litellm_params)
|
||||
def _deployment_litellm_params(deployment: Mapping[str, object] | None) -> Mapping[str, object]:
|
||||
return as_str_mapping((deployment or {}).get("litellm_params")) or MappingProxyType({})
|
||||
|
||||
|
||||
def _claude_count_translator(route: tuple[str, str] | None) -> AnthropicConfig:
|
||||
if route is None:
|
||||
return AnthropicConfig()
|
||||
model, provider = route
|
||||
|
|
@ -88,12 +92,14 @@ def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicC
|
|||
return chat_config if isinstance(chat_config, AnthropicConfig) else AnthropicConfig()
|
||||
|
||||
|
||||
def _claude_options_provider(litellm_params: Mapping[str, object], use_mantle: bool) -> LlmProviders:
|
||||
def _claude_options_provider(route: tuple[str, str] | None, use_mantle: bool) -> LlmProviders:
|
||||
if not use_mantle:
|
||||
return LlmProviders.BEDROCK
|
||||
route: Final = _deployment_route(litellm_params)
|
||||
deployment_provider: Final = None if route is None else route[1]
|
||||
return LlmProviders.BEDROCK if deployment_provider == LlmProviders.BEDROCK.value else LlmProviders.BEDROCK_MANTLE
|
||||
if route is None:
|
||||
return LlmProviders.BEDROCK_MANTLE
|
||||
model, provider = route
|
||||
runs_on_invoke: Final = provider == LlmProviders.BEDROCK.value and not _runs_on_converse(model)
|
||||
return LlmProviders.BEDROCK if runs_on_invoke else LlmProviders.BEDROCK_MANTLE
|
||||
|
||||
|
||||
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
|
||||
|
|
@ -259,7 +265,8 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
counted_model: Final = model_to_use.removeprefix("bedrock_mantle/")
|
||||
resolved_model: Final = get_bedrock_base_model(counted_model)
|
||||
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE
|
||||
litellm_params: Final = (deployment or {}).get("litellm_params", {})
|
||||
litellm_params: Final = _deployment_litellm_params(deployment)
|
||||
route: Final = _deployment_route(litellm_params) if is_mantle_claude_model(resolved_model) else None
|
||||
translated: Final = (
|
||||
await anthropic_count_fields(
|
||||
counted_model,
|
||||
|
|
@ -267,7 +274,7 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
system,
|
||||
tools,
|
||||
messages_format,
|
||||
_claude_count_translator(litellm_params),
|
||||
_claude_count_translator(route),
|
||||
chat_params,
|
||||
)
|
||||
if is_mantle_claude_model(resolved_model)
|
||||
|
|
@ -280,7 +287,7 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
claude_options: Final = (
|
||||
claude_count_options(
|
||||
counted_model,
|
||||
_claude_options_provider(litellm_params, use_mantle),
|
||||
_claude_options_provider(route, use_mantle),
|
||||
sent_messages,
|
||||
sent_tools,
|
||||
_options_converse_sends(request_options, deployment_inference_model(deployment)),
|
||||
|
|
@ -322,7 +329,7 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
else await self._count_with(
|
||||
self._runtime_handler.handle_count_tokens_request(
|
||||
request_data=dict(runtime_request),
|
||||
litellm_params=litellm_params,
|
||||
litellm_params=dict(litellm_params),
|
||||
resolved_model=resolved_model,
|
||||
client=self._client,
|
||||
input_type="invokeModel" if runtime_fields else None,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import base64
|
||||
import functools
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
|
|
@ -13,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_EFFORT_BETA_HEADER,
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
|
||||
ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,
|
||||
)
|
||||
from litellm.types.utils import CountContextManagement
|
||||
|
|
@ -56,6 +59,13 @@ def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def aws_credentials_in_the_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret")
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() -> None:
|
||||
bedrock: Final = _Bedrock(runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))
|
||||
|
|
@ -476,6 +486,7 @@ async def _messages_inference_and_count(
|
|||
tool: Mapping[str, JsonValue] = _WEATHER_TOOL,
|
||||
runtime: tuple[int, Mapping[str, object]] = (200, {"inputTokens": 412}),
|
||||
inference_reply: Mapping[str, object] = _CLAUDE_REPLY,
|
||||
anthropic_beta: str | None = None,
|
||||
) -> tuple[_Bedrock, _Bedrock]:
|
||||
credentials: Final = DEPLOYMENT["litellm_params"]
|
||||
inference: Final = _Bedrock(runtime=(200, inference_reply), mantle=(200, inference_reply))
|
||||
|
|
@ -484,6 +495,7 @@ async def _messages_inference_and_count(
|
|||
messages=MESSAGES,
|
||||
tools=[dict(tool)],
|
||||
max_tokens=4096,
|
||||
extra_headers=None if anthropic_beta is None else {"anthropic-beta": anthropic_beta},
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)),
|
||||
**request_options,
|
||||
**credentials,
|
||||
|
|
@ -502,6 +514,7 @@ async def _messages_inference_and_count(
|
|||
else CountContextManagement(spec=request_options["context_management"], drop_invalid_edits=False)
|
||||
),
|
||||
request_options=request_options,
|
||||
anthropic_beta=anthropic_beta,
|
||||
)
|
||||
assert result is not None
|
||||
assert result.error is False
|
||||
|
|
@ -677,7 +690,7 @@ _CALLER_BETA_CASES: Final = pytest.mark.parametrize(
|
|||
)
|
||||
|
||||
|
||||
async def _invoke_inference_betas(caller_beta: str, model_info: dict[str, JsonValue]) -> tuple[JsonValue, ...]:
|
||||
async def _invoke_inference_betas(caller_beta: str, model_info: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]:
|
||||
inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY))
|
||||
await litellm.anthropic_messages(
|
||||
model=_PER_TURN_CONTROL_DEPLOYMENT_MODEL,
|
||||
|
|
@ -692,22 +705,22 @@ async def _invoke_inference_betas(caller_beta: str, model_info: dict[str, JsonVa
|
|||
return _betas_besides_effort(_JSON_BODY.validate_json(inference.requests[0].content))
|
||||
|
||||
|
||||
def _per_turn_control_deployment(model_info: dict[str, JsonValue]) -> dict[str, object]:
|
||||
def _per_turn_control_deployment(model_info: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"litellm_params": {**DEPLOYMENT["litellm_params"], "model": _PER_TURN_CONTROL_DEPLOYMENT_MODEL},
|
||||
"model_info": model_info,
|
||||
"model_info": dict(model_info),
|
||||
}
|
||||
|
||||
|
||||
async def _count_with_caller_beta(
|
||||
model_to_use: str, caller_beta: str, deployment: dict[str, object] | None
|
||||
model_to_use: str, caller_beta: str, deployment: Mapping[str, JsonValue] | None
|
||||
) -> _Bedrock:
|
||||
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177}))
|
||||
result: Final = await _counter(bedrock).count_tokens(
|
||||
model_to_use=model_to_use,
|
||||
messages=MESSAGES,
|
||||
contents=None,
|
||||
deployment=deployment,
|
||||
deployment=None if deployment is None else dict(deployment),
|
||||
request_model="claude",
|
||||
request_options=_THINKING,
|
||||
anthropic_beta=caller_beta,
|
||||
|
|
@ -720,7 +733,7 @@ async def _count_with_caller_beta(
|
|||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
@_CALLER_BETA_CASES
|
||||
async def test_runtime_count_keeps_and_drops_the_caller_betas_invoke_inference_keeps_and_drops(
|
||||
caller_beta: str, model_info: dict[str, JsonValue], kept: bool
|
||||
caller_beta: str, model_info: Mapping[str, JsonValue], kept: bool
|
||||
) -> None:
|
||||
sent: Final = await _invoke_inference_betas(caller_beta, model_info)
|
||||
bedrock: Final = await _count_with_caller_beta(
|
||||
|
|
@ -737,7 +750,7 @@ async def test_runtime_count_keeps_and_drops_the_caller_betas_invoke_inference_k
|
|||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
@_CALLER_BETA_CASES
|
||||
async def test_mantle_count_for_a_bedrock_deployment_keeps_and_drops_the_caller_betas_invoke_inference_keeps_and_drops(
|
||||
caller_beta: str, model_info: dict[str, JsonValue], kept: bool
|
||||
caller_beta: str, model_info: Mapping[str, JsonValue], kept: bool
|
||||
) -> None:
|
||||
sent: Final = await _invoke_inference_betas(caller_beta, model_info)
|
||||
bedrock: Final = await _count_with_caller_beta(
|
||||
|
|
@ -750,22 +763,65 @@ async def test_mantle_count_for_a_bedrock_deployment_keeps_and_drops_the_caller_
|
|||
assert (caller_beta in counted) is kept
|
||||
|
||||
|
||||
_INTERLEAVED_THINKING: Final = "interleaved-thinking-2025-05-14"
|
||||
_MANTLE_COUNT_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
async def test_runtime_count_without_a_deployment_drops_the_caller_betas_invoke_inference_drops(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret")
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
interleaved_thinking: Final = "interleaved-thinking-2025-05-14"
|
||||
sent: Final = await _invoke_inference_betas(interleaved_thinking, {})
|
||||
bedrock: Final = await _count_with_caller_beta("us.anthropic.claude-opus-5-5", interleaved_thinking, None)
|
||||
@pytest.mark.usefixtures("local_beta_headers_config", "aws_credentials_in_the_environment")
|
||||
async def test_runtime_count_without_a_deployment_drops_the_caller_betas_invoke_inference_drops() -> None:
|
||||
sent: Final = await _invoke_inference_betas(_INTERLEAVED_THINKING, {})
|
||||
bedrock: Final = await _count_with_caller_beta("us.anthropic.claude-opus-5-5", _INTERLEAVED_THINKING, None)
|
||||
|
||||
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
|
||||
counted: Final = _betas_besides_effort(_invoke_body(bedrock.requests[0]))
|
||||
assert counted == sent
|
||||
assert interleaved_thinking not in counted
|
||||
|
||||
|
||||
async def _mantle_inference_betas(caller_beta: str, inference_params: Mapping[str, str]) -> frozenset[str]:
|
||||
inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY))
|
||||
await litellm.anthropic_messages(
|
||||
messages=MESSAGES,
|
||||
max_tokens=4096,
|
||||
extra_headers={"anthropic-beta": caller_beta},
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)),
|
||||
**_THINKING,
|
||||
**DEPLOYMENT["litellm_params"],
|
||||
**inference_params,
|
||||
)
|
||||
assert inference.posted_hosts() == (MANTLE_HOST,)
|
||||
return frozenset(_sent_betas(inference.requests[0]))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_beta_headers_config", "aws_credentials_in_the_environment")
|
||||
@pytest.mark.parametrize(
|
||||
"deployment_params",
|
||||
(
|
||||
{"model": _MANTLE_COUNT_MODEL},
|
||||
{"model": _MANTLE_COUNT_MODEL.removeprefix("bedrock_mantle/"), "custom_llm_provider": "bedrock_mantle"},
|
||||
None,
|
||||
),
|
||||
ids=("bedrock-mantle-deployment", "bedrock-mantle-by-custom-llm-provider", "no-deployment"),
|
||||
)
|
||||
async def test_mantle_count_for_a_mantle_deployment_or_none_sends_the_caller_betas_mantle_inference_sends(
|
||||
deployment_params: Mapping[str, str] | None,
|
||||
) -> None:
|
||||
sent: Final = await _mantle_inference_betas(
|
||||
_INTERLEAVED_THINKING, deployment_params or {"model": _MANTLE_COUNT_MODEL}
|
||||
)
|
||||
invoke_sent: Final = frozenset(await _invoke_inference_betas(_INTERLEAVED_THINKING, {}))
|
||||
bedrock: Final = await _count_with_caller_beta(
|
||||
_MANTLE_COUNT_MODEL,
|
||||
_INTERLEAVED_THINKING,
|
||||
None
|
||||
if deployment_params is None
|
||||
else {"litellm_params": {**DEPLOYMENT["litellm_params"], **deployment_params}},
|
||||
)
|
||||
|
||||
assert bedrock.posted_hosts() == (MANTLE_HOST,)
|
||||
assert sent != invoke_sent
|
||||
assert frozenset(_sent_betas(bedrock.requests[0])) == sent
|
||||
|
||||
|
||||
_CONVERSE_REPLY: Final = {
|
||||
|
|
@ -806,6 +862,45 @@ async def test_runtime_count_sends_disable_parallel_tool_use_only_when_messages_
|
|||
)
|
||||
|
||||
|
||||
def _at(value: JsonValue, *path: str) -> JsonValue:
|
||||
return functools.reduce(lambda node, key: node.get(key) if isinstance(node, dict) else None, path, value)
|
||||
|
||||
|
||||
def _content_texts(message: JsonValue) -> tuple[str, ...]:
|
||||
content: Final = message.get("content") if isinstance(message, dict) else None
|
||||
return (content,) if isinstance(content, str) else _texts(message)
|
||||
|
||||
|
||||
def _message_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
messages: Final = body.get("messages")
|
||||
assert isinstance(messages, list)
|
||||
return tuple(itertools.chain.from_iterable(map(_content_texts, messages)))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas_converse_inference_sends() -> None:
|
||||
inference, bedrock = await _messages_inference_and_count(
|
||||
"bedrock/converse/us.anthropic.claude-opus-5-5",
|
||||
_MANTLE_COUNT_MODEL,
|
||||
{**_CITY_FORMAT, **_THINKING},
|
||||
inference_reply=_CONVERSE_REPLY,
|
||||
anthropic_beta=f"{ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER},context-1m-2025-08-07",
|
||||
)
|
||||
|
||||
assert (inference.posted_hosts(), bedrock.posted_hosts()) == ((RUNTIME_HOST,), (MANTLE_HOST,))
|
||||
sent: Final = _JSON_BODY.validate_json(inference.requests[0].content)
|
||||
counted: Final = _JSON_BODY.validate_json(bedrock.requests[0].content)
|
||||
sent_schema: Final = _at(sent, "outputConfig", "textFormat", "structure", "jsonSchema", "schema")
|
||||
assert isinstance(sent_schema, str)
|
||||
assert _at(counted, "output_config", "format") == {
|
||||
"type": _at(sent, "outputConfig", "textFormat", "type"),
|
||||
"schema": json.loads(sent_schema),
|
||||
}
|
||||
assert _message_texts(counted) == _message_texts(sent)
|
||||
assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_on_a_stand_in_model_counts_the_messages_and_output_config_the_deployment_model_sends() -> None:
|
||||
inference, bedrock = await _messages_inference_and_count(
|
||||
|
|
|
|||
|
|
@ -1967,16 +1967,19 @@ _BEDROCK_CLAUDE_DEPLOYMENT: Final[dict[str, str]] = {
|
|||
"aws_secret_access_key": "test-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
_THINKING_BODY: Final[dict[str, JsonValue]] = {"thinking": {"type": "enabled", "budget_tokens": 2000}}
|
||||
_EFFORT_BODY: Final[dict[str, JsonValue]] = {"output_config": {"effort": "low"}}
|
||||
_DROPS_THE_EFFORT: Final[dict[str, JsonValue]] = {"additional_drop_params": ["output_config.effort"]}
|
||||
_ANTHROPIC_DEPLOYMENT: Final[dict[str, JsonValue]] = {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-ant-test"}
|
||||
_AZURE_AI_DEPLOYMENT: Final[dict[str, JsonValue]] = {
|
||||
_THINKING_BODY: Final[Mapping[str, JsonValue]] = {"thinking": {"type": "enabled", "budget_tokens": 2000}}
|
||||
_EFFORT_BODY: Final[Mapping[str, JsonValue]] = {"output_config": {"effort": "low"}}
|
||||
_DROPS_THE_EFFORT: Final[Mapping[str, JsonValue]] = {"additional_drop_params": ["output_config.effort"]}
|
||||
_ANTHROPIC_HOST: Final = "api.anthropic.com"
|
||||
_ANTHROPIC_DEPLOYMENT: Final[Mapping[str, str]] = {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-ant-test"}
|
||||
_AZURE_AI_HOST: Final = "my-resource.services.ai.azure.com"
|
||||
_AZURE_AI_DEPLOYMENT: Final[Mapping[str, str]] = {
|
||||
"model": "azure_ai/claude-haiku-5-5",
|
||||
"api_base": "https://my-resource.services.ai.azure.com/anthropic",
|
||||
"api_base": f"https://{_AZURE_AI_HOST}/anthropic",
|
||||
"api_key": "sk-azure-test",
|
||||
}
|
||||
_VERTEX_DEPLOYMENT: Final[dict[str, JsonValue]] = {
|
||||
_VERTEX_HOST: Final = "us-east5-aiplatform.googleapis.com"
|
||||
_VERTEX_DEPLOYMENT: Final[Mapping[str, str]] = {
|
||||
"model": "vertex_ai/claude-haiku-5-5",
|
||||
"vertex_project": "test-project",
|
||||
"vertex_location": "us-east5",
|
||||
|
|
@ -1985,17 +1988,32 @@ _VERTEX_DEPLOYMENT: Final[dict[str, JsonValue]] = {
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_params", "body", "caller_beta", "provider_calls"),
|
||||
("deployment_params", "body", "caller_beta", "counted_hosts"),
|
||||
(
|
||||
({"model": "gemini/gemini-3.8-flash", "api_key": "sk-gemini-test"}, _THINKING_BODY, None, 1),
|
||||
(_ANTHROPIC_DEPLOYMENT, _THINKING_BODY, None, 2),
|
||||
({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1),
|
||||
({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, _CALLER_BETA, 2),
|
||||
(_AZURE_AI_DEPLOYMENT, _THINKING_BODY, None, 2),
|
||||
({**_AZURE_AI_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1),
|
||||
(_VERTEX_DEPLOYMENT, _THINKING_BODY, None, 2),
|
||||
({**_VERTEX_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1),
|
||||
({**_BEDROCK_CLAUDE_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 2),
|
||||
(
|
||||
{"model": "gemini/gemini-3.8-flash", "api_key": "sk-gemini-test"},
|
||||
_THINKING_BODY,
|
||||
None,
|
||||
("generativelanguage.googleapis.com",),
|
||||
),
|
||||
(_ANTHROPIC_DEPLOYMENT, _THINKING_BODY, None, (_ANTHROPIC_HOST, _ANTHROPIC_HOST)),
|
||||
({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_ANTHROPIC_HOST,)),
|
||||
(
|
||||
{**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT},
|
||||
_EFFORT_BODY,
|
||||
_CALLER_BETA,
|
||||
(_ANTHROPIC_HOST, _ANTHROPIC_HOST),
|
||||
),
|
||||
(_AZURE_AI_DEPLOYMENT, _THINKING_BODY, None, (_AZURE_AI_HOST, _AZURE_AI_HOST)),
|
||||
({**_AZURE_AI_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_AZURE_AI_HOST,)),
|
||||
(_VERTEX_DEPLOYMENT, _THINKING_BODY, None, (_VERTEX_HOST, _VERTEX_HOST)),
|
||||
({**_VERTEX_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_VERTEX_HOST,)),
|
||||
(
|
||||
{**_BEDROCK_CLAUDE_DEPLOYMENT, **_DROPS_THE_EFFORT},
|
||||
_EFFORT_BODY,
|
||||
None,
|
||||
(_BEDROCK_RUNTIME_HOST, _BEDROCK_MANTLE_HOST),
|
||||
),
|
||||
),
|
||||
ids=(
|
||||
"counter-without-claude-options",
|
||||
|
|
@ -2011,10 +2029,10 @@ _VERTEX_DEPLOYMENT: Final[dict[str, JsonValue]] = {
|
|||
)
|
||||
async def test_anthropic_count_tokens_counts_again_without_the_options_only_when_the_counter_forwarded_them(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
deployment_params: dict[str, JsonValue],
|
||||
body: dict[str, JsonValue],
|
||||
deployment_params: Mapping[str, JsonValue],
|
||||
body: Mapping[str, JsonValue],
|
||||
caller_beta: str | None,
|
||||
provider_calls: int,
|
||||
counted_hosts: tuple[str, ...],
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
|
|
@ -2028,7 +2046,7 @@ async def test_anthropic_count_tokens_counts_again_without_the_options_only_when
|
|||
"python-httpx/0.28.1",
|
||||
{} if caller_beta is None else {"anthropic-beta": caller_beta},
|
||||
)
|
||||
router: Final = Router(model_list=[{"model_name": "claude-agent", "litellm_params": deployment_params}])
|
||||
router: Final = Router(model_list=[{"model_name": "claude-agent", "litellm_params": dict(deployment_params)}])
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig())
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
|
@ -2048,7 +2066,7 @@ async def test_anthropic_count_tokens_counts_again_without_the_options_only_when
|
|||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 400
|
||||
assert route.call_count == provider_calls
|
||||
assert tuple(call.request.url.host for call in route.calls) == counted_hosts
|
||||
|
||||
|
||||
async def _count_on_a_bedrock_claude_deployment(
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ This test validates:
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from queue import SimpleQueue
|
||||
from typing import Dict, Final, List, get_args
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -798,8 +799,8 @@ async def test_direct_anthropic_deployment_sends_the_caller_betas_its_row_marks_
|
|||
|
||||
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
|
||||
row: Final[dict[str, str | None]] = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()["anthropic"]
|
||||
unsupported: Final = min(beta for beta, mapped in row.items() if mapped is None and "-" in beta)
|
||||
row: Final[Mapping[str, str | None]] = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()["anthropic"]
|
||||
unsupported: Final = min(beta for beta, mapped in row.items() if mapped is None)
|
||||
sent: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
|
||||
def reply(request: httpx.Request) -> httpx.Response:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue