From af3d5c15189953fcac6c74723e4dc6dadd575c8d Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 16 Oct 2026 02:15:04 +0000 Subject: [PATCH] refactor: call execute_chat from the Messages to chat bridge The /v1/messages bridge for chat-only models now plans the chat call once and runs execute_chat directly instead of re-entering litellm.acompletion, so deployment hooks run once per request. Requests that need the chat surface (fallbacks, model_list, MCP gateway tools, prompt management, litellm_system_prompt) still go through litellm.acompletion --- .../litellm_core_utils/streaming_handler.py | 2 +- .../pass_through/adapters/handler.py | 13 +- litellm/main.py | 200 +++++++++++++++++ litellm/utils.py | 6 +- .../adapters/test_handler_chat_executor.py | 204 ++++++++++++++++++ .../context_management/test_constants.py | 19 +- ...erimental_pass_through_messages_handler.py | 14 +- ...test_vertex_and_google_ai_studio_gemini.py | 7 +- 8 files changed, 440 insertions(+), 25 deletions(-) create mode 100644 tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..19bc4b66fd7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1702,7 +1702,7 @@ class CustomStreamWrapper: setattr(e, "message", str(e)) raise self._map_stream_exception(e) - def set_logging_event_loop(self, loop): + def set_logging_event_loop(self, loop: asyncio.AbstractEventLoop) -> None: """ import litellm, asyncio diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 01e781485a0..d869136d8e2 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -39,6 +39,7 @@ from litellm.llms.anthropic.pass_through.utils import ( proxy_sub_call_logging_kwargs, proxy_sub_call_trace_id, ) +from litellm.main import aexecute_planned_chat, execute_planned_chat, plan_chat from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -796,8 +797,13 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response: Final = await litellm.acompletion(**completion_kwargs) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) + planned_chat: Final = None if logging_obj is None else plan_chat(completion_kwargs, logging_obj) + completion_response: Final = ( + await litellm.acompletion(**completion_kwargs) + if planned_chat is None + else await aexecute_planned_chat(planned_chat) + ) if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( @@ -942,8 +948,11 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response: Final = litellm.completion(**completion_kwargs) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) + planned_chat: Final = None if logging_obj is None else plan_chat(completion_kwargs, logging_obj) + completion_response: Final = ( + litellm.completion(**completion_kwargs) if planned_chat is None else execute_planned_chat(planned_chat) + ) if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( diff --git a/litellm/main.py b/litellm/main.py index af6ddf361ea..8cbea971cac 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -23,6 +23,7 @@ from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mappin from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy +from dataclasses import dataclass from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -6496,6 +6497,205 @@ def execute_chat( ) +_PLANNED_CHAT_ARGS: Final = frozenset( + ("kwargs", "non_default_params", "control_options", "skip_responses_api_bridge", "prompt_id", "prompt_variables") +) +_EXECUTE_CHAT_CODE: Final = execute_chat.__code__ +_EXECUTE_CHAT_PARAMS: Final = ( + frozenset(_EXECUTE_CHAT_CODE.co_varnames[: _EXECUTE_CHAT_CODE.co_kwonlyargcount]) - _PLANNED_CHAT_ARGS +) +_CHAT_SURFACE_KWARGS: Final = ( + "fallbacks", + "model_list", + "prompt_id", + "litellm_system_prompt", + "ensure_alternating_roles", + "user_continue_message", + "assistant_continue_message", + "mock_delay", + "mock_timeout", +) +_TEXT_COMPLETION_PROVIDERS: Final = frozenset( + ("text-completion-openai", "text-completion-codestral", "text-completion-inception") +) + + +@dataclass(frozen=True, slots=True) +class PlannedChat: + model: str + custom_llm_provider: str + named: Mapping[str, object] + kwargs: Mapping[str, object] + non_default_params: Mapping[str, object] + control_options: ControlOptions + + +def _optional_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + +_OBJECT_LIST: Final = TypeAdapter(list[object]) + + +def _object_list(value: object) -> list[object] | None: # mutable-ok: feeds validate_and_fix_openai_* which take a list + return _OBJECT_LIST.validate_python(value) if isinstance(value, list) else None + + +def _tool_choice(value: object) -> str | dict[str, object] | None: # mutable-ok: its validator takes a dict + if isinstance(value, str): + return value + return _OPTIONAL_PARAMS_ADAPTER.validate_python(value) if isinstance(value, dict) else None + + +def _stop(value: object) -> str | list[str] | None: # mutable-ok: validate_openai_optional_params takes a list + if isinstance(value, list): + return [item for item in _OBJECT_LIST.validate_python(value) if isinstance(item, str)] + return value if isinstance(value, str) else None + + +def _thinking(value: object) -> AnthropicThinkingParam | bool | None: + if isinstance(value, dict): + return cast(AnthropicThinkingParam, value) # cast-ok: the caller's thinking dict, normalized by the validator + return value if isinstance(value, bool) else None + + +def _chat_entry_points_wrapped() -> bool: + """Tracers such as ddtrace wrap `litellm.completion`/`acompletion`, so bridged calls must go through them.""" + from litellm.chat_completions import dispatch + + return litellm.completion is not dispatch.completion or litellm.acompletion is not dispatch.acompletion + + +def _uses_chat_surface(request: Mapping[str, object]) -> bool: + from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + from litellm.types.llms.openai import ToolParam + + tools: Final = request.get("tools") + return ( + litellm.model_fallbacks is not None + or _chat_entry_points_wrapped() + or any(request.get(key) is not None for key in _CHAT_SURFACE_KWARGS) + or ( + isinstance(tools, list) + and LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway( + tools=cast(Iterable[ToolParam], tools) # cast-ok: OpenAI-format tools from the caller's request + ) + ) + ) + + +def plan_chat(request: Mapping[str, object], logging_obj: LiteLLMLoggingObj) -> PlannedChat | None: + model: Final = request.get("model") + if not isinstance(model, str) or _uses_chat_surface(request): + return None + from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook + + custom_llm_provider: Final = _optional_str(request.get("custom_llm_provider")) + api_base: Final = _optional_str(request.get("api_base")) + _, resolved_provider, _, _ = get_llm_provider( + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + ) + if resolved_provider in _TEXT_COMPLETION_PROVIDERS: + return None + kwargs: Final = {key: value for key, value in request.items() if key not in _EXECUTE_CHAT_PARAMS} + tools: Final = validate_and_fix_openai_tools(tools=_object_list(request.get("tools"))) + messages: Final = get_completion_messages( + messages=validate_and_fix_openai_messages(messages=_object_list(request.get("messages")) or []), + ensure_alternating_roles=False, + user_continue_message=None, + assistant_continue_message=None, + ) + non_default_params: Final = get_non_default_completion_params(kwargs=kwargs) + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=non_default_params, + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + tools=tools, + enable_prompt_caching=request.get("enable_prompt_caching") is True, + api_base=api_base, + ) + if logging_obj.should_run_prompt_management_hooks(non_default_params=non_default_params, tools=tools): + return None + named: Final = {key: value for key, value in request.items() if key in _EXECUTE_CHAT_PARAMS} + tool_choice: Final = validate_chat_completion_tool_choice( + tool_choice=_tool_choice(request.get("tool_choice")), model=model + ) + return PlannedChat( + model=model, + custom_llm_provider=resolved_provider, + named=MappingProxyType( + { + **named, + "messages": messages, + "tools": tools, + "tool_choice": tool_choice, + "stop": validate_openai_optional_params(stop=_stop(request.get("stop"))), + "thinking": validate_and_fix_thinking_param(thinking=_thinking(request.get("thinking"))), + } + ), + kwargs=MappingProxyType(kwargs), + non_default_params=MappingProxyType(non_default_params), + control_options=_resolve_control_options(kwargs, model), + ) + + +def _run_planned_chat(plan: PlannedChat, is_async: bool) -> ModelResponse | CustomStreamWrapper: + return execute_chat( + **plan.named, # pyright: ignore[reportArgumentType] # named holds the caller's validated OpenAI-format params + kwargs={**plan.kwargs, "acompletion": is_async}, + non_default_params=dict(plan.non_default_params), + control_options=plan.control_options, + skip_responses_api_bridge=False, + prompt_id=None, + prompt_variables=None, + ) + + +def execute_planned_chat(plan: PlannedChat) -> ModelResponse | CustomStreamWrapper: + return _run_planned_chat(plan, is_async=False) + + +async def _settled_chat_response( + plan: PlannedChat, loop: asyncio.AbstractEventLoop +) -> ModelResponse | CustomStreamWrapper: + pending: Final[object] = await loop.run_in_executor( + None, partial(contextvars.copy_context().run, partial(_run_planned_chat, plan, True)) + ) + if asyncio.iscoroutine(pending): + return await _resolve_dispatched_chat_response(pending) + if isinstance(pending, (ModelResponse, CustomStreamWrapper)): + return pending + raise TypeError(f"chat provider returned {type(pending).__name__}, expected a ModelResponse or a stream") + + +async def _with_agentic_loop(plan: PlannedChat, response: ModelResponse) -> ModelResponse | CustomStreamWrapper: + looped: Final = await maybe_run_chat_completion_agentic_loop( + response=response, + model=plan.model, + messages=_object_list(plan.named["messages"]) or [], + optional_params={ + key: value + for key, value in plan.named.items() + if value is not None and key not in ("model", "messages", "stream", "deployment_id") + }, + kwargs=dict(plan.kwargs), + logging_obj=plan.kwargs.get("litellm_logging_obj"), + custom_llm_provider=plan.custom_llm_provider, + stream=plan.named.get("stream") is True, + ) + return response if looped is None else looped + + +async def aexecute_planned_chat(plan: PlannedChat) -> ModelResponse | CustomStreamWrapper: + loop: Final = asyncio.get_event_loop() + response: Final = await _settled_chat_response(plan, loop) + if isinstance(response, CustomStreamWrapper): + response.set_logging_event_loop(loop=loop) + return response + return await _with_agentic_loop(plan, response) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries diff --git a/litellm/utils.py b/litellm/utils.py index 6b16fc398bd..23823a8f354 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8839,7 +8839,7 @@ def validate_and_fix_openai_messages(messages: list): return validate_chat_completion_user_messages(messages=new_messages) -def validate_and_fix_openai_tools(tools: list | None) -> list[dict] | None: +def validate_and_fix_openai_tools(tools: list | None) -> list[dict[str, object]] | None: """ Ensure tools is List[dict] and not List[BaseModel] """ @@ -8931,9 +8931,9 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( - tool_choice: dict | str | None, + tool_choice: dict[str, object] | str | None, model: str = "", -) -> dict | str | None: +) -> dict[str, object] | str | None: """ Confirm the tool choice is passed in the OpenAI format. diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py new file mode 100644 index 00000000000..39351cbacfd --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_chat_executor.py @@ -0,0 +1,204 @@ +import datetime +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, TypeAlias + +import httpx +import pytest +import respx + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.pass_through.adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from litellm.types.llms.openai import AllMessageValues +from litellm.types.prompts.init_prompts import PromptSpec +from litellm.types.utils import StandardCallbackDynamicParams + +API_BASE: Final = "http://chat-upstream.test/v1" +CHAT_RESPONSE: Final = { + "id": "chatcmpl-bridge", + "object": "chat.completion", + "created": 1, + "model": "served-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "from upstream"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, +} + + +class _DeploymentHookCounter(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.pre_call = 0 + self.post_call = 0 + + async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: object) -> None: + self.pre_call += 1 + + async def async_post_call_success_deployment_hook( + self, request_data: Mapping[str, object], response: object, call_type: object + ) -> None: + self.post_call += 1 + + +_PromptResult: TypeAlias = tuple[str, list[AllMessageValues], dict[str, object]] + + +class _TemplatePromptManager(CustomPromptManagement): + async def async_get_chat_completion_prompt( + self, + model: str, + messages: Sequence[AllMessageValues], + non_default_params: Mapping[str, object], + prompt_id: str | None, + prompt_variables: Mapping[str, object] | None, + dynamic_callback_params: StandardCallbackDynamicParams, + litellm_logging_obj: Logging, + prompt_spec: PromptSpec | None = None, + tools: Sequence[Mapping[str, object]] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> _PromptResult: + return model, [{"role": "user", "content": "from the template"}], {**non_default_params} + + +@pytest.fixture +def upstream(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(f"{API_BASE}/chat/completions").mock(return_value=httpx.Response(200, json=CHAT_RESPONSE)) + + +async def _ask(extra: Mapping[str, object] = MappingProxyType({})) -> object: + return await litellm.anthropic.messages.acreate( + model="hosted_vllm/served-model", + api_base=API_BASE, + api_key="sk-chat", + max_tokens=10, + messages=[{"role": "user", "content": "hi"}], + **extra, + ) + + +@pytest.mark.asyncio +async def test_bridge_runs_deployment_hooks_once_per_request(upstream: respx.Route, monkeypatch: pytest.MonkeyPatch): + counter: Final = _DeploymentHookCounter() + monkeypatch.setattr(litellm, "callbacks", [counter]) + + response: Final = await _ask() + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "from upstream"}] + assert upstream.call_count == 1 + assert (counter.pre_call, counter.post_call) == (1, 1) + + +@pytest.mark.asyncio +async def test_bridge_streams_the_upstream_reply(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + chunk: Final = '{"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"streamed"},"finish_reason":null}]}' + done: Final = '{"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + respx_mock.post(f"{API_BASE}/chat/completions").mock( + return_value=httpx.Response( + 200, + content=f"data: {chunk}\n\ndata: {done}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + ) + ) + + stream: Final = await _ask({"stream": True}) + body: Final = b"".join([event async for event in stream]).decode() + + assert "event: message_start" in body + assert '"text":"streamed"' in body.replace(" ", "") + + +@pytest.mark.asyncio +async def test_bridge_applies_the_prompt_template_named_by_prompt_id( + upstream: respx.Route, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "callbacks", [_TemplatePromptManager()]) + + await _ask({"prompt_id": "bridge_template"}) + + assert "from the template" in upstream.calls.last.request.content.decode() + + +@pytest.mark.asyncio +async def test_bridge_maps_an_upstream_429_to_rate_limit_error( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{API_BASE}/chat/completions").mock( + return_value=httpx.Response(429, json={"error": {"message": "slow down", "type": "rate_limit"}}) + ) + + with pytest.raises(litellm.RateLimitError): + await _ask() + + +def test_sync_bridge_answers_from_one_upstream_call(upstream: respx.Route): + logging_obj: Final = Logging( + model="hosted_vllm/served-model", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="6825beef-0000-4000-8000-000000000005", + function_id="1234", + ) + + response: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=10, + messages=[{"role": "user", "content": "hi"}], + model="hosted_vllm/served-model", + api_base=API_BASE, + api_key="sk-chat", + litellm_logging_obj=logging_obj, + ) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "from upstream"}] + assert upstream.call_count == 1 + + +@pytest.mark.asyncio +async def test_bridge_runs_prompt_management_triggered_without_prompt_id( + upstream: respx.Route, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "callbacks", [_TemplatePromptManager()]) + + await _ask({"vector_store_ids": ["vs_bridge"]}) + + assert "from the template" in upstream.calls.last.request.content.decode() + + +@pytest.mark.asyncio +async def test_bridge_keeps_the_litellm_system_prompt(upstream: respx.Route): + await _ask({"litellm_system_prompt": "answer like a pirate"}) + + assert "answer like a pirate" in upstream.calls.last.request.content.decode() + + +@pytest.mark.asyncio +async def test_bridge_goes_through_a_wrapped_acompletion(upstream: respx.Route, monkeypatch: pytest.MonkeyPatch): + unwrapped: Final = litellm.acompletion + seen: Final[list[object]] = [] # mutable-ok: records the calls the wrapper sees + + async def traced_acompletion(*args: object, **kwargs: object) -> object: + seen.append(kwargs.get("model")) + return await unwrapped(*args, **kwargs) + + monkeypatch.setattr(litellm, "acompletion", traced_acompletion) + + response: Final = await _ask() + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "from upstream"}] + assert seen == ["hosted_vllm/served-model"] + assert upstream.call_count == 1 diff --git a/tests/unit/llms/anthropic/pass_through/context_management/test_constants.py b/tests/unit/llms/anthropic/pass_through/context_management/test_constants.py index 680c2b6f79a..2808f71780c 100644 --- a/tests/unit/llms/anthropic/pass_through/context_management/test_constants.py +++ b/tests/unit/llms/anthropic/pass_through/context_management/test_constants.py @@ -103,15 +103,18 @@ async def _mock_streaming_chunks(): ) +_PLANNED_CHAT = "litellm.llms.anthropic.pass_through.adapters.handler.aexecute_planned_chat" + + @pytest.mark.asyncio async def test_polyfill_round_trip_non_streaming(): captured = {} - async def fake_acompletion(**kwargs): - captured.update(kwargs) + async def fake_chat(plan): + captured.update({**plan.kwargs, **plan.named}) return _mock_completion_response() - with patch("litellm.acompletion", side_effect=fake_acompletion): + with patch(_PLANNED_CHAT, side_effect=fake_chat): response = await litellm.anthropic.messages.acreate( model=MODEL, messages=_make_history(n_pairs=5), @@ -173,11 +176,11 @@ async def test_polyfill_round_trip_non_streaming(): async def test_polyfill_trigger_not_met_passes_through_unchanged(): captured = {} - async def fake_acompletion(**kwargs): - captured.update(kwargs) + async def fake_chat(plan): + captured.update({**plan.kwargs, **plan.named}) return _mock_completion_response() - with patch("litellm.acompletion", side_effect=fake_acompletion): + with patch(_PLANNED_CHAT, side_effect=fake_chat): response = await litellm.anthropic.messages.acreate( model=MODEL, messages=_make_history(n_pairs=2), @@ -209,10 +212,10 @@ async def test_polyfill_trigger_not_met_passes_through_unchanged(): @pytest.mark.asyncio async def test_polyfill_streaming_attaches_to_message_delta(): - async def fake_acompletion(**kwargs): + async def fake_chat(plan): return _mock_streaming_chunks() - with patch("litellm.acompletion", side_effect=fake_acompletion): + with patch(_PLANNED_CHAT, side_effect=fake_chat): response = await litellm.anthropic.messages.acreate( model=MODEL, messages=_make_history(n_pairs=5), diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 1f02b3aaa5a..fc4bfce9b4d 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -286,8 +286,8 @@ async def test_bedrock_converse_budget_tokens_preserved(): usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ) - with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: - mock_acompletion.return_value = mock_response + with patch("litellm.llms.anthropic.pass_through.adapters.handler.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat: + mock_chat.return_value = mock_response try: await messages.acreate( @@ -299,13 +299,9 @@ async def test_bedrock_converse_budget_tokens_preserved(): except (ValueError, TypeError, AttributeError): pass # Expected due to response format conversion - mock_acompletion.assert_called_once() + mock_chat.assert_called_once() - call_kwargs = mock_acompletion.call_args.kwargs - print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) - - # Verify thinking parameter is passed through with budget_tokens preserved - thinking_param = call_kwargs.get("thinking") + thinking_param = mock_chat.call_args.args[0].named.get("thinking") assert thinking_param is not None, "thinking parameter should be passed to acompletion" assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" assert thinking_param.get("budget_tokens") == 1024, ( @@ -1187,7 +1183,7 @@ async def test_messages_streaming_reports_provider_local_model(requested_model, async def __anext__(self): raise StopAsyncIteration - with patch("litellm.acompletion", new=AsyncMock(return_value=_EmptyStream())): + with patch("litellm.llms.anthropic.pass_through.adapters.handler.aexecute_planned_chat", new=AsyncMock(return_value=_EmptyStream())): stream = await litellm.anthropic.messages.acreate( max_tokens=100, messages=[{"role": "user", "content": "ping"}], diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 636d4034ccc..97dbf47da9e 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -7783,8 +7783,11 @@ async def test_claude_tool_use_with_gemini(): return chunk raise StopAsyncIteration - with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: - mock_acompletion.return_value = MockAsyncStream() + with patch( + "litellm.llms.anthropic.pass_through.adapters.handler.aexecute_planned_chat", + new_callable=AsyncMock, + ) as mock_chat: + mock_chat.return_value = MockAsyncStream() response = await litellm.anthropic.messages.acreate( messages=[