mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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
This commit is contained in:
Yuneng Jiang
parent
1aba849af2
commit
859b72fe96
1 changed files with 115 additions and 83 deletions
|
|
@ -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"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue