diff --git a/litellm/fusion_router.py b/litellm/fusion_router.py index 1139a915080..48882e95c21 100644 --- a/litellm/fusion_router.py +++ b/litellm/fusion_router.py @@ -36,6 +36,24 @@ FUSION_PROTOCOL_VERSION: Final = "fusion-tool-v1" _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _OBJECT_MAPPINGS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _BUDGET_RESERVATION_METADATA_KEY: Final = "user_api_key_budget_reservation" +_RESPONSES_ONLY_REQUEST_KEYS: Final = frozenset( + { + "background", + "include", + "input", + "instructions", + "max_output_tokens", + "max_tool_calls", + "partial_images", + "previous_response_id", + "prompt", + "prompt_cache_options", + "reasoning", + "stream_options", + "text", + "truncation", + } +) def is_fusion_router_model(model: str) -> bool: @@ -283,7 +301,9 @@ def _internal_kwargs( kwargs = { key: value for key, value in request_kwargs.items() - if key not in _INTERNAL_REQUEST_KEYS and key not in _INTERNAL_RESPONSE_KEYS + if key not in _INTERNAL_REQUEST_KEYS + and key not in _INTERNAL_RESPONSE_KEYS + and key not in _RESPONSES_ONLY_REQUEST_KEYS } kwargs.pop("metadata", None) kwargs.pop("litellm_metadata", None) @@ -350,6 +370,60 @@ def _fusion_tool_call(response: ModelResponse) -> ChatCompletionMessageToolCall return None +def _mixed_tool_call_indexes(response: ModelResponse) -> tuple[frozenset[int], tuple[int, ...]]: + if not response.choices: + return frozenset(), () + tool_calls = response.choices[0].message.tool_calls or () + fusion_indexes = frozenset( + index + for index, tool_call in enumerate(tool_calls) + if isinstance(tool_call, ChatCompletionMessageToolCall) and tool_call.function.name == FUSION_TOOL_NAME + ) + client_indexes = tuple(index for index in range(len(tool_calls)) if index not in fusion_indexes) + return fusion_indexes, client_indexes + + +def _without_mixed_fusion_tool_call(response: ModelResponse) -> tuple[ModelResponse, frozenset[int]]: + """Prefer executable client calls when a provider violates the one-path contract.""" + fusion_indexes, client_indexes = _mixed_tool_call_indexes(response) + if not fusion_indexes or not client_indexes: + return response, frozenset() + sanitized = response.model_copy(deep=True) + tool_calls = sanitized.choices[0].message.tool_calls or () + sanitized.choices[0].message.tool_calls = [tool_calls[index] for index in client_indexes] + return sanitized, fusion_indexes + + +def _without_stream_tool_call_indexes( + chunks: Sequence[ModelResponseStream], + removed_indexes: frozenset[int], +) -> list[ModelResponseStream]: + if not removed_indexes: + return list(chunks) + kept_indexes = sorted( + { + tool_call.index + for chunk in chunks + for choice in chunk.choices + for tool_call in (choice.delta.tool_calls or ()) + if tool_call.index not in removed_indexes + } + ) + index_map = {old_index: new_index for new_index, old_index in enumerate(kept_indexes)} + sanitized_chunks: list[ModelResponseStream] = [] + for chunk in chunks: + sanitized = chunk.model_copy(deep=True) + for choice in sanitized.choices: + tool_calls = choice.delta.tool_calls or () + choice.delta.tool_calls = [ + tool_call.model_copy(update={"index": index_map[tool_call.index]}) + for tool_call in tool_calls + if tool_call.index in index_map + ] or None + sanitized_chunks.append(sanitized) + return sanitized_chunks + + def _fusion_query(tool_call: ChatCompletionMessageToolCall) -> str | None: try: arguments = _OBJECT_MAPPING_ADAPTER.validate_json(tool_call.function.arguments) @@ -584,7 +658,8 @@ def _outer_kwargs(request_kwargs: Mapping[str, object]) -> dict[str, object]: return { key: value for key, value in request_kwargs.items() - if key + if key not in _RESPONSES_ONLY_REQUEST_KEYS + and key not in frozenset( { "_fusion_depth", @@ -671,6 +746,7 @@ class FusionRouter: # from spend reconciliation. litellm_metadata=metadata, max_tokens_per_page=1024, + _fusion_proxy_auth_required=isinstance(request_kwargs.get("proxy_server_request"), Mapping), ) if isinstance(result, BaseModel): result = result.model_dump() @@ -754,7 +830,8 @@ class FusionRouter: **kwargs, ) if isinstance(response, ModelResponse): - return response, None + sanitized_response, _ = _without_mixed_fusion_tool_call(response) + return sanitized_response, None chunks: list[ModelResponseStream] = [] try: @@ -774,9 +851,10 @@ class FusionRouter: llm_provider="", model=self.config.outer_model, ) + built, removed_indexes = _without_mixed_fusion_tool_call(built) replay = FusionReplayStream( source=response, - chunks=chunks, + chunks=_without_stream_tool_call_indexes(chunks, removed_indexes), fusion_metadata={"invoked": False, "protocol": FUSION_PROTOCOL_VERSION}, ) return built, replay diff --git a/litellm/router.py b/litellm/router.py index edc5a501eda..3e0705236d3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9468,7 +9468,7 @@ class Router: ) -> object: """Late-bound Search API bridge with the originating caller's permissions.""" metadata_values: Final = tuple(kwargs.get(key) for key in ("litellm_metadata", "metadata")) - user_api_key_auth: Final = next( + raw_user_api_key_auth: Final = next( ( metadata.get("user_api_key_auth") for metadata in metadata_values @@ -9476,17 +9476,32 @@ class Router: ), None, ) - if user_api_key_auth is not None: - from litellm.proxy._types import UserAPIKeyAuth + # Direct Router usage has no proxy identity. The Fusion caller sets this + # private flag only when the originating request came through the proxy. + proxy_auth_required: Final = kwargs.pop("_fusion_proxy_auth_required", False) is True + if raw_user_api_key_auth is not None or proxy_auth_required: + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.search_endpoints.endpoints import ( authorize_search_tool_call, ) - if isinstance(user_api_key_auth, UserAPIKeyAuth): - await authorize_search_tool_call( - search_tool_name=model, - user_api_key_dict=user_api_key_auth, + try: + user_api_key_auth = ( + raw_user_api_key_auth + if isinstance(raw_user_api_key_auth, UserAPIKeyAuth) + else UserAPIKeyAuth.model_validate(raw_user_api_key_auth) ) + except ValidationError as exc: + raise ProxyException( + message="Fusion Search Tool authorization context is missing or invalid", + type=ProxyErrorTypes.auth_error, + param=None, + code=403, + ) from exc + await authorize_search_tool_call( + search_tool_name=model, + user_api_key_dict=user_api_key_auth, + ) return await self.asearch(model=model, query=query, **kwargs) def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: diff --git a/tests/test_litellm/test_fusion_router.py b/tests/test_litellm/test_fusion_router.py index 6b7ee7a8cbc..7bb0f593970 100644 --- a/tests/test_litellm/test_fusion_router.py +++ b/tests/test_litellm/test_fusion_router.py @@ -12,12 +12,14 @@ import litellm from litellm.fusion_router import ( FUSION_TOOL_NAME, FusionRouterConfig, + _without_stream_tool_call_indexes, build_fusion_router, fusion_router_dependencies, validate_fusion_router_write, ) from litellm.router import Router from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponseStream from litellm.utils import CustomStreamWrapper, ModelResponse @@ -163,6 +165,99 @@ async def test_outer_client_tool_call_is_returned_without_running_panel_or_secon assert [call["model"] for call in completion.calls] == ["outer"] +@pytest.mark.asyncio +async def test_mixed_fusion_and_client_tool_calls_return_only_executable_client_calls() -> None: + client_call = { + "id": "email-1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to":"user@example.com"}'}, + } + mixed_response = _response( + None, + [ + { + "id": "fusion-call-1", + "type": "function", + "function": {"name": FUSION_TOOL_NAME, "arguments": '{"query":"Investigate this"}'}, + }, + client_call, + ], + ) + completion = RecordingCompletion({"outer": [mixed_response]}) + + response = await _router(completion).acompletion( + messages=[{"role": "user", "content": "Research and send the update"}], + stream=False, + request_kwargs={ + "tools": [ + { + "type": "function", + "function": {"name": "send_email", "parameters": {"type": "object"}}, + } + ] + }, + ) + + assert isinstance(response, ModelResponse) + assert [call.function.name for call in response.choices[0].message.tool_calls] == ["send_email"] + assert [call["model"] for call in completion.calls] == ["outer"] + assert response._hidden_params["fusion"]["invoked"] is False + + +def test_mixed_stream_removes_private_call_and_reindexes_client_call() -> None: + chunks = [ + ModelResponseStream( + choices=[ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "fusion-call-1", + "type": "function", + "function": {"name": FUSION_TOOL_NAME, "arguments": '{"query":"test"}'}, + }, + { + "index": 1, + "id": "email-1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to":"user@example.com"}'}, + }, + ] + } + } + ] + ) + ] + + sanitized = _without_stream_tool_call_indexes(chunks, frozenset({0})) + + tool_calls = sanitized[0].choices[0].delta.tool_calls + assert len(tool_calls) == 1 + assert tool_calls[0].index == 0 + assert tool_calls[0].function.name == "send_email" + + +@pytest.mark.asyncio +async def test_responses_only_kwargs_never_reach_fusion_chat_calls() -> None: + completion = RecordingCompletion({"outer": [_response("Final")]}) + + await _router(completion).acompletion( + messages=[{"role": "system", "content": "Follow instructions"}, {"role": "user", "content": "Answer"}], + stream=False, + request_kwargs={ + "input": "raw Responses input", + "instructions": "Follow instructions", + "previous_response_id": "resp-1", + "include": ["reasoning.encrypted_content"], + "text": {"format": {"type": "text"}}, + }, + ) + + outer_call = completion.calls[0] + assert not {"input", "instructions", "previous_response_id", "include", "text"} & outer_call.keys() + + @pytest.mark.asyncio async def test_forced_fusion_runs_parallel_panel_then_analyst_then_outer() -> None: completion = RecordingCompletion( @@ -415,12 +510,16 @@ async def test_configured_search_tool_is_private_to_panel_and_analyst() -> None: response = await _router(completion, search=search, search_tool_name="web-search", max_tool_calls=4).acompletion( messages=[{"role": "user", "content": "Research this"}], stream=False, - request_kwargs={"litellm_metadata": {"user_api_key_budget_reservation": reservation}}, + request_kwargs={ + "litellm_metadata": {"user_api_key_budget_reservation": reservation}, + "proxy_server_request": {"body": {}}, + }, ) assert isinstance(response, ModelResponse) assert search_calls[0]["model"] == "web-search" assert search_calls[0]["query"] == "current evidence" + assert search_calls[0]["_fusion_proxy_auth_required"] is True assert search_calls[0]["litellm_metadata"]["internal_call_origin"] == "fusion_research" assert search_calls[0]["litellm_metadata"]["user_api_key_budget_reservation"] is reservation second_panel_call = [call for call in completion.calls if call["model"] == "panel-a"][1] @@ -637,13 +736,32 @@ async def test_fusion_search_checks_proxy_permissions_before_router_search(monke await router._fusion_asearch( # pyright: ignore[reportPrivateUsage] model="restricted-search", query="evidence", - litellm_metadata={"user_api_key_auth": user_api_key_auth}, + litellm_metadata={"user_api_key_auth": user_api_key_auth.model_dump()}, ) get_team_object.assert_awaited_once() raw_search.assert_not_awaited() +@pytest.mark.asyncio +async def test_fusion_search_fails_closed_when_proxy_auth_context_is_missing(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._types import ProxyException + + router = Router(model_list=[]) + raw_search = AsyncMock(return_value={"results": []}) + monkeypatch.setattr(router, "asearch", raw_search) + + with pytest.raises(ProxyException, match="authorization context is missing or invalid"): + await router._fusion_asearch( # pyright: ignore[reportPrivateUsage] + model="restricted-search", + query="evidence", + litellm_metadata={}, + _fusion_proxy_auth_required=True, + ) + + raw_search.assert_not_awaited() + + @pytest.mark.asyncio async def test_router_responses_and_anthropic_adapters_use_same_fusion_model() -> None: router = Router(model_list=_router_model_list())