refactor: expose MCP chained-call helpers

This commit is contained in:
tarunag10 2026-08-20 13:28:50 +05:30
parent 887d8c751d
commit 7768fe9b00
4 changed files with 9 additions and 7 deletions

View file

@ -453,7 +453,7 @@ async def acompletion_with_mcp(
)
# Make follow-up call with streaming
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params(
self.base_call_args,
original_stream_setting=True,
)
@ -628,7 +628,7 @@ async def acompletion_with_mcp(
)
# Make follow-up call with original stream setting
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params(
base_call_args,
original_stream_setting=stream,
)

View file

@ -72,7 +72,7 @@ class LiteLLM_Proxy_MCP_Handler:
"""
@staticmethod
def _prepare_chained_call_params(
def prepare_chained_call_params(
params: Mapping[str, Any],
) -> dict[str, Any]: # mutable-ok: returns a sanitized copy for the next provider call
"""Copy request params without state owned by the previous LLM call.
@ -1261,14 +1261,14 @@ class LiteLLM_Proxy_MCP_Handler:
return initial_params
@staticmethod
def _prepare_follow_up_call_params(call_params: dict[str, Any], original_stream_setting: bool) -> dict[str, Any]:
def prepare_follow_up_call_params(call_params: dict[str, Any], original_stream_setting: bool) -> dict[str, Any]:
"""
Prepare call parameters for the follow-up LLM call after tool execution.
Restores the original streaming setting and removes tool_choice since
we're now providing tool results, not requesting tool calls.
"""
follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(call_params)
follow_up_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_chained_call_params(call_params)
# Restore original streaming setting for follow-up call
follow_up_params["stream"] = original_stream_setting
@ -1278,6 +1278,8 @@ class LiteLLM_Proxy_MCP_Handler:
return follow_up_params
_prepare_follow_up_call_params = prepare_follow_up_call_params
@staticmethod
def _add_mcp_output_elements_to_response(
response: ResponsesAPIResponse,

View file

@ -782,7 +782,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
)
# Make follow-up call with streaming
follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(
follow_up_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_chained_call_params(
self.original_request_params
)
follow_up_params.update(

View file

@ -459,7 +459,7 @@ def test_prepare_follow_up_call_params_resets_per_call_logging_state():
},
}
follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original, original_stream_setting=True)
follow_up = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params(original, original_stream_setting=True)
assert "litellm_call_id" not in follow_up
assert "litellm_logging_obj" not in follow_up