From 9d435f09bf4174ba78d6345f332fb690483ed58e Mon Sep 17 00:00:00 2001 From: Yug Date: Wed, 29 Apr 2026 18:14:50 +0530 Subject: [PATCH] resolve --- .../mcp_server/sampling_handler.py | 2 +- tests/mcp_tests/test_sampling_coverage.py | 57 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 tests/mcp_tests/test_sampling_coverage.py diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index e600c300b59..590268077ea 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -503,7 +503,7 @@ async def handle_sampling_create_message( response = await llm_router.acompletion(**completion_kwargs) else: response = await litellm.acompletion(**completion_kwargs) - except Exception: + except ImportError: response = await litellm.acompletion(**completion_kwargs) # 7. Convert response to MCP format result = _convert_openai_response_to_mcp_result( diff --git a/tests/mcp_tests/test_sampling_coverage.py b/tests/mcp_tests/test_sampling_coverage.py new file mode 100644 index 00000000000..e0a9a5af19d --- /dev/null +++ b/tests/mcp_tests/test_sampling_coverage.py @@ -0,0 +1,57 @@ +from mcp.types import TextContent, ImageContent +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_single_content, +) + + +def test_convert_single_content_coverage(): + # text content + txt = TextContent(type="text", text="hello") + res = _convert_single_content(txt) + assert res == {"type": "text", "text": "hello"} + + # image content + img = ImageContent(type="image", data="base64", mimeType="image/png") + res_img = _convert_single_content(img) + assert res_img == { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,base64"}, + } + + +def test_convert_mcp_tool_choice_coverage(): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tool_choice_to_openai, + ) + + class MockToolChoice: + def __init__(self, mode): + self.mode = mode + + assert _convert_mcp_tool_choice_to_openai(MockToolChoice("auto")) == "auto" + assert _convert_mcp_tool_choice_to_openai(MockToolChoice("required")) == "required" + assert _convert_mcp_tool_choice_to_openai(MockToolChoice("none")) == "none" + + +def test_convert_openai_response_to_mcp_result_coverage(): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_openai_response_to_mcp_result, + ) + from litellm import ModelResponse, Message, Choices + + # Text response + mock_resp = ModelResponse( + id="test-id", + choices=[ + Choices( + message=Message(content="hello", role="assistant"), finish_reason="stop" + ) + ], + model="gpt-4o", + ) + mcp_res = _convert_openai_response_to_mcp_result(mock_resp, "gpt-4o") + assert mcp_res.role == "assistant" + assert mcp_res.content.type == "text" + assert mcp_res.content.text == "hello" + assert mcp_res.model == "gpt-4o" + assert mcp_res.stopReason == "endTurn"