scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(router): preserve fusion tool boundaries

Browse source
This commit is contained in:
moe-berri 2026-09-03 13:50:08 -07:00
parent 4859ac43da
commit 2852fb846a
3 changed files with 224 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

86
litellm/fusion_router.py
View file

@ -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

29
litellm/router.py
View file

@ -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:

122
tests/test_litellm/test_fusion_router.py
View file

@ -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())