mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
refactor: expose MCP chained-call helpers
This commit is contained in:
parent
887d8c751d
commit
7768fe9b00
4 changed files with 9 additions and 7 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue