From 859b72fe96bb1ccf1e7b752b44590de72923e049 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Jul 2026 19:03:29 -0700 Subject: [PATCH] test(llm_translation): assert real behavior in the streamed tool-response test test_function_calling_with_tool_response had no assertions. It looped up to three live streaming rounds, printed the result, and swallowed ServiceUnavailableError, so it passed no matter what the provider returned. Its only real failure mode was running past CI's 120s cap, and because the job uses --timeout_method=thread that kills the worker with os._exit, which xdist reports as "worker 'gwN' crashed" The test now checks what its name claims: the streamed argument deltas reassemble into valid JSON, the tool call carries an id, the assistant tool_calls turn round-trips back into messages, and the follow-up answer is grounded in the tool result. Each of the two rounds carries an explicit 30s request timeout, so a stalled stream raises litellm.Timeout and skips instead of blowing the pytest cap and taking the worker down with it Also drops the litellm._turn_on_debug() call, which had no teardown and left DEBUG logging on for every later test on that worker, and switches to self.completion_function so the Router subclass actually exercises the router here Verified by mutation: making the provider transform drop the tool message content is caught by the grounded-answer assertion (the old test passed that mutation), and truncating each streamed argument delta is caught by the JSON reassembly check --- tests/llm_translation/base_llm_unit_tests.py | 198 +++++++++++-------- 1 file changed, 115 insertions(+), 83 deletions(-) diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index 6d845f4b2f1..0ebeb4b4368 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -2,7 +2,7 @@ import httpx import json import pytest import sys -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from unittest.mock import MagicMock, Mock, patch import os from litellm._uuid import uuid @@ -65,6 +65,30 @@ def _usage_format_tests(usage: litellm.Usage): assert usage.prompt_tokens > usage.prompt_tokens_details.cached_tokens +_TOOL_CALL_STREAM_TIMEOUT = 30 + + +def _drain_tool_call_stream(stream) -> Tuple[str, List[Any]]: + """ + Consume a streaming completion, returning the concatenated text content and the + tool calls reassembled from their per-index argument deltas. + """ + content = "" + tool_calls_by_index: Dict[Any, Any] = {} + for chunk in stream: + delta = chunk.choices[0].delta + if delta.content: + content += delta.content + for tool_call in delta.tool_calls or []: + accumulated = tool_calls_by_index.get(tool_call.index) + if accumulated is None: + tool_call.function.arguments = tool_call.function.arguments or "" + tool_calls_by_index[tool_call.index] = tool_call + else: + accumulated.function.arguments += tool_call.function.arguments or "" + return content, list(tool_calls_by_index.values()) + + class BaseLLMChatTest(ABC): """ Abstract base test class that enforces a common test across all test classes. @@ -1332,102 +1356,110 @@ class BaseLLMChatTest(ABC): ), "Audio URL not sent to gemini" def test_function_calling_with_tool_response(self): + """ + A streamed tool call must reassemble into valid JSON arguments, round-trip back + into `messages` as an assistant turn, and let the model answer from the tool result. + """ from litellm.utils import supports_function_calling - from litellm import completion - litellm._turn_on_debug() - try: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") + base_completion_call_args = self.get_base_completion_call_args() + if not supports_function_calling(base_completion_call_args["model"], None): + pytest.skip("Model does not support function calling") - base_completion_call_args = self.get_base_completion_call_args() - if not supports_function_calling(base_completion_call_args["model"], None): - print("Model does not support function calling") - pytest.skip("Model does not support function calling") - - def get_weather(city: str): - return f"City: {city}, Weather: Sunny with 34 degree Celcius" - - TOOLS = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the weather in a city", - "parameters": { - "$id": "https://some/internal/name", - "$schema": "https://json-schema.org/draft-07/schema", - "type": "object", - "properties": { - "city": { - "type": "string", - "description": "The city to get the weather for", - } - }, - "required": ["city"], - "additionalProperties": False, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather in a city", + "parameters": { + "$id": "https://some/internal/name", + "$schema": "https://json-schema.org/draft-07/schema", + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "The city to get the weather for", + } }, - "strict": True, + "required": ["city"], + "additionalProperties": False, }, - } - ] + "strict": True, + }, + } + ] + messages: List[Dict[str, Any]] = [ + {"content": "How is the weather in Mumbai?", "role": "user"} + ] - messages = [{"content": "How is the weather in Mumbai?", "role": "user"}] - response, iteration = "", 0 - while True: - if response: - break - # Create a streaming response with tool calling enabled - stream = completion( + try: + content, tool_calls = _drain_tool_call_stream( + self.completion_function( **base_completion_call_args, messages=messages, - tools=TOOLS, + tools=tools, stream=True, + timeout=_TOOL_CALL_STREAM_TIMEOUT, ) + ) - final_tool_calls = {} - for chunk in stream: - delta = chunk.choices[0].delta - print(delta) - if delta.content: - response += delta.content - elif delta.tool_calls: - for tool_call in chunk.choices[0].delta.tool_calls or []: - index = tool_call.index - if index not in final_tool_calls: - final_tool_calls[index] = tool_call - else: - final_tool_calls[ - index - ].function.arguments += tool_call.function.arguments - if final_tool_calls: - for tool_call in final_tool_calls.values(): - if tool_call.function.name == "get_weather": - city = json.loads(tool_call.function.arguments)["city"] - tool_response = get_weather(city) - messages.append( - { - "role": "assistant", - "tool_calls": [tool_call], - "content": None, - } - ) - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "content": tool_response, - } - ) - iteration += 1 - if iteration > 2: - print("Something went wrong!") - break + assert [ + tool_call.function.name for tool_call in tool_calls + ] == ["get_weather"], ( + f"expected a single streamed get_weather call, got {tool_calls} " + f"with content={content!r}" + ) - print(response) + tool_call = tool_calls[0] + assert tool_call.id, "streamed tool call has no id to reply to" + + raw_arguments = tool_call.function.arguments + try: + arguments = json.loads(raw_arguments) + except json.JSONDecodeError as e: + pytest.fail( + f"streamed argument deltas did not reassemble into valid JSON: " + f"{raw_arguments!r} ({e})" + ) + assert "city" in arguments, f"tool call arguments missing 'city': {arguments}" + assert "mumbai" in arguments["city"].lower() + + messages.append( + {"role": "assistant", "tool_calls": [tool_call], "content": None} + ) + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": f"City: {arguments['city']}, Weather: Sunny with 34 degree Celcius", + } + ) + + final_content, final_tool_calls = _drain_tool_call_stream( + self.completion_function( + **base_completion_call_args, + messages=messages, + tools=tools, + stream=True, + timeout=_TOOL_CALL_STREAM_TIMEOUT, + ) + ) + + assert final_content.strip(), ( + "model returned no answer after being given the tool response " + f"(tool_calls={final_tool_calls})" + ) + assert ( + "34" in final_content or "sunny" in final_content.lower() + ), f"answer is not grounded in the tool response: {final_content!r}" + except litellm.Timeout: + pytest.skip("Model took too long to respond") except litellm.ServiceUnavailableError: - pass + pytest.skip("Model is unavailable") def test_reasoning_effort(self): """Test that reasoning_effort is passed correctly to the model"""