diff --git a/cookbook/fusion_models.md b/cookbook/fusion_models.md index 4c9f42e26d8..0a95b1bd770 100644 --- a/cookbook/fusion_models.md +++ b/cookbook/fusion_models.md @@ -41,6 +41,8 @@ Call `fusion/general` exactly like any other model. `invocation: auto` lets the `reasoning_effort: none` makes deliberation replace private extended reasoning where a provider supports that parameter. LiteLLM drops it for providers that do not support it. The optional Search Tool supplies search results and bounded page content through LiteLLM's Search API. `max_tool_calls` limits searches separately for every panel model and the analyst; it defaults to 4 and is capped at 16. This first version does not expose a separate URL-fetch tool. +`panel_timeout_seconds` bounds each panel member and the complete analyst phase, including any Search Tool loop. If the analyst times out, the outer model still receives the successful raw panel responses. + The outer model must support function calling. Panel and analyst models only need function calling when a Search Tool is configured. Granting access to the Fusion model lets the request use its administrator-configured model and search dependencies; the panel query and private research are sent to those deployments under their normal provider data policies. ## Operational behavior diff --git a/litellm/fusion_router.py b/litellm/fusion_router.py index 1f2ddb2594e..c47563490c5 100644 --- a/litellm/fusion_router.py +++ b/litellm/fusion_router.py @@ -365,8 +365,7 @@ def _research_tool_calls(response: ModelResponse) -> tuple[ChatCompletionMessage return tuple( tool_call for tool_call in response.choices[0].message.tool_calls or () - if isinstance(tool_call, ChatCompletionMessageToolCall) - and tool_call.function.name == "litellm_fusion_search" + if isinstance(tool_call, ChatCompletionMessageToolCall) and tool_call.function.name == "litellm_fusion_search" ) @@ -680,13 +679,9 @@ class FusionRouter: selected_calls = search_calls[:remaining_searches] if not selected_calls: return response - current_messages.append( - cast(AllMessageValues, response.choices[0].message.model_dump(exclude_none=True)) - ) + current_messages.append(cast(AllMessageValues, response.choices[0].message.model_dump(exclude_none=True))) current_messages.extend( - await asyncio.gather( - *(self._execute_research_call(call, request_kwargs) for call in selected_calls) - ) + await asyncio.gather(*(self._execute_research_call(call, request_kwargs) for call in selected_calls)) ) current_messages.extend( { @@ -761,7 +756,7 @@ class FusionRouter: if self.config.reasoning_effort is not None: kwargs["reasoning_effort"] = self.config.reasoning_effort try: - response = await asyncio.wait_for( + response: Final = await asyncio.wait_for( self._call_internal_model( model=model, messages=panel_messages, @@ -802,11 +797,16 @@ class FusionRouter: if self.config.reasoning_effort is not None: kwargs["reasoning_effort"] = self.config.reasoning_effort try: - response = await self._call_internal_model( - model=model, - messages=messages, - kwargs=kwargs, - request_kwargs=request_kwargs, + # One timeout bounds the complete private analyst phase, including + # any configured Search Tool loop, just as it bounds panel members. + response: Final = await asyncio.wait_for( + self._call_internal_model( + model=model, + messages=messages, + kwargs=kwargs, + request_kwargs=request_kwargs, + ), + timeout=self.config.panel_timeout_seconds, ) except Exception: return None @@ -840,6 +840,12 @@ class FusionRouter: hidden["fusion"] = fusion_metadata return replay_stream if replay_stream is not None else initial_response + # A streamed initial response is fully buffered to discover the private + # Fusion call. The direct-response path returns the replay wrapper to + # the caller; the invocation path suppresses it, so it owns cleanup. + if replay_stream is not None: + await replay_stream.aclose() + fusion_metadata["invoked"] = True raw_query = _fusion_query(tool_call) if raw_query is None: diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 2676682c59d..0574e095d90 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -14,6 +14,46 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin router: Final = APIRouter() +async def authorize_search_tool_call( + search_tool_name: str, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Apply the proxy's key and team Search Tool allowlists. + + Keep this check reusable by internal features that call ``Router.asearch`` + directly, because those calls do not pass through the HTTP search endpoint. + """ + from litellm.proxy.auth.auth_checks import ( + can_key_call_search_tool, + can_team_call_search_tool, + get_team_object, + ) + + await can_key_call_search_tool( + search_tool_name=search_tool_name, + valid_token=user_api_key_dict, + ) + + if user_api_key_dict.team_id: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team_object: Final = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + await can_team_call_search_tool( + search_tool_name=search_tool_name, + team_object=team_object, + ) + + @router.post( "/v1/search/{search_tool_name}", dependencies=[Depends(user_api_key_auth)], @@ -138,39 +178,11 @@ async def search( data["model"] = data["search_tool_name"] search_tool_name_value: Final = data["search_tool_name"] - # Authorization check: verify key can access this search tool - from litellm.proxy.auth.auth_checks import ( - can_key_call_search_tool, - can_team_call_search_tool, - get_team_object, - ) - try: - # Check key-level access - await can_key_call_search_tool( + await authorize_search_tool_call( search_tool_name=search_tool_name_value, - valid_token=user_api_key_dict, + user_api_key_dict=user_api_key_dict, ) - - # Check team-level access if key is associated with a team - if user_api_key_dict.team_id: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - team_object: Final = await get_team_object( - team_id=user_api_key_dict.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - await can_team_call_search_tool( - search_tool_name=search_tool_name_value, - team_object=team_object, - ) except Exception as e: verbose_proxy_logger.error("Search tool authorization failed for %s: %s", search_tool_name_value, e) raise diff --git a/litellm/router.py b/litellm/router.py index 9e88b432452..edc5a501eda 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9463,8 +9463,30 @@ class Router: search=self._fusion_asearch, ) - async def _fusion_asearch(self, *, model: str, query: str, **kwargs: object) -> object: - """Late-bound Search API bridge; Fusion routers are registered before endpoint factories run.""" + async def _fusion_asearch( # kwargs-ok: bridge preserves the Router.asearch keyword surface + self, *, model: str, query: str, **kwargs: object + ) -> 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( + ( + metadata.get("user_api_key_auth") + for metadata in metadata_values + if isinstance(metadata, Mapping) and metadata.get("user_api_key_auth") is not None + ), + None, + ) + if user_api_key_auth is not None: + from litellm.proxy._types import 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, + ) 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 35f8666b785..38ee2710f4f 100644 --- a/tests/test_litellm/test_fusion_router.py +++ b/tests/test_litellm/test_fusion_router.py @@ -4,6 +4,7 @@ import json from collections import deque from collections.abc import Mapping, Sequence from typing import Final +from unittest.mock import AsyncMock import pytest @@ -270,6 +271,58 @@ async def test_partial_panel_and_invalid_analyst_degrade_to_raw_responses() -> N ] +@pytest.mark.asyncio +async def test_analyst_timeout_degrades_to_raw_panel_responses() -> None: + class HangingAnalystCompletion(RecordingCompletion): + def __init__(self) -> None: + super().__init__( + { + "outer": [_fusion_call(), _response("Final")], + "panel-a": [_response("Panel A")], + "panel-b": [_response("Panel B")], + } + ) + self.analyst_started = asyncio.Event() + self.analyst_cancelled = asyncio.Event() + + async def __call__( + self, + *, + model: str, + messages: list[AllMessageValues], + stream: bool, + **kwargs: object, + ) -> ModelResponse | CustomStreamWrapper: + if model != "analyst": + return await super().__call__(model=model, messages=messages, stream=stream, **kwargs) + self.calls.append({"model": model, "messages": messages, "stream": stream, **kwargs}) + self.analyst_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + self.analyst_cancelled.set() + raise + raise AssertionError("unreachable") + + completion = HangingAnalystCompletion() + response = await asyncio.wait_for( + _router(completion, panel_timeout_seconds=0.2).acompletion( + messages=[{"role": "user", "content": "Hard question"}], + stream=False, + request_kwargs={}, + ), + timeout=2, + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Final" + assert completion.analyst_started.is_set() + assert completion.analyst_cancelled.is_set() + assert response._hidden_params["fusion"]["analysis_available"] is False + payload = json.loads(completion.calls[-1]["messages"][-1]["content"]) + assert [item["content"] for item in payload["responses"]] == ["Panel A", "Panel B"] + + @pytest.mark.asyncio async def test_all_panel_failures_are_a_typed_tool_result_the_outer_can_recover_from() -> None: completion = RecordingCompletion( @@ -479,6 +532,65 @@ async def test_router_replays_direct_outer_response_as_an_async_stream() -> None assert response._hidden_params["fusion"]["invoked"] is False +@pytest.mark.asyncio +async def test_invoked_fusion_closes_suppressed_initial_stream(monkeypatch: pytest.MonkeyPatch) -> None: + completion = RecordingCompletion( + { + "outer": [_response("Final")], + "panel-a": [_response("Panel A")], + "panel-b": [_response("Panel B")], + "analyst": [_response(_analysis())], + } + ) + router = _router(completion) + replay_stream = AsyncMock() + monkeypatch.setattr( + router, + "_initial_outer_call", + AsyncMock(return_value=(_fusion_call(), replay_stream)), + ) + + response = await router.acompletion( + messages=[{"role": "user", "content": "Answer"}], + stream=True, + request_kwargs={}, + ) + + assert isinstance(response, ModelResponse) + replay_stream.aclose.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_fusion_search_checks_proxy_permissions_before_router_search(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.models.object_permission import LiteLLM_ObjectPermissionTable + from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth + from litellm.proxy.auth import auth_checks + + router = Router(model_list=[]) + user_api_key_auth = UserAPIKeyAuth(team_id="restricted-team") + team = LiteLLM_TeamTable( + team_id="restricted-team", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="team-permissions", + search_tools=["allowed-search"], + ), + ) + get_team_object = AsyncMock(return_value=team) + raw_search = AsyncMock(return_value={"results": []}) + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(router, "asearch", raw_search) + + with pytest.raises(ProxyException, match="Team not allowed to access search tool"): + await router._fusion_asearch( # pyright: ignore[reportPrivateUsage] + model="restricted-search", + query="evidence", + litellm_metadata={"user_api_key_auth": user_api_key_auth}, + ) + + get_team_object.assert_awaited_once() + 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()) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/FusionModels/FusionModelsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/FusionModels/FusionModelsPanel.tsx index fa118852da8..3299524235e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/FusionModels/FusionModelsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/FusionModels/FusionModelsPanel.tsx @@ -315,7 +315,7 @@ function FusionModelDialog({ {advancedOpen && (
- + 600) { - return "Panel timeout must be between 1 and 600 seconds."; + return "Panel and analyst timeout must be between 1 and 600 seconds."; } if (value.max_candidate_chars < 1000 || value.max_candidate_chars > 50000) { return "Candidate limit must be between 1,000 and 50,000 characters.";