From c6e4d372cd8ef5732a7a8b2a75c423bcf0aa70bc Mon Sep 17 00:00:00 2001 From: Rajasimhareddybolla Date: Sun, 25 Oct 2026 18:26:23 +0530 Subject: [PATCH] fix(types): resolve basedpyright diagnostics and add alternating incremental multi-tool streaming test --- .../litellm_core_utils/streaming_handler.py | 25 +- litellm/llms/databricks/streaming_utils.py | 6 +- .../chat/test_openai_like_chat_handler.py | 383 ++++++++++++++++++ 3 files changed, 397 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 32e6e97570f..92bc0e9e6a7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1366,13 +1366,9 @@ class CustomStreamWrapper: raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] - if ( - "reasoning_content" in anthropic_response_obj - and anthropic_response_obj["reasoning_content"] is not None - ): - completion_obj["reasoning_content"] = anthropic_response_obj[ # rebind-ok: forwards reasoning - "reasoning_content" - ] + reasoning_chunk: Final = anthropic_response_obj.get("reasoning_content") + if reasoning_chunk is not None: + completion_obj["reasoning_content"] = reasoning_chunk # rebind-ok: forwards reasoning if anthropic_response_obj["is_finished"]: self.received_finish_reason = anthropic_response_obj["finish_reason"] @@ -1386,10 +1382,12 @@ class CustomStreamWrapper: litellm.Usage(**anthropic_response_obj["usage"]), ) - if "tool_calls" in anthropic_response_obj and anthropic_response_obj["tool_calls"] is not None: - completion_obj["tool_calls"] = anthropic_response_obj["tool_calls"] # rebind-ok: forwards tool calls - elif "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None: - completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] # rebind-ok: forwards tool calls + tool_calls_chunk: Final = anthropic_response_obj.get("tool_calls") + legacy_tool_chunk: Final = anthropic_response_obj.get("tool_use") + if tool_calls_chunk is not None: + completion_obj["tool_calls"] = tool_calls_chunk # rebind-ok: forwards tool calls + elif legacy_tool_chunk is not None: + completion_obj["tool_calls"] = [legacy_tool_chunk] # rebind-ok: forwards tool calls if ( "provider_specific_fields" in anthropic_response_obj @@ -2798,9 +2796,8 @@ def convert_generic_chunk_to_model_response_stream( from litellm.types.utils import Delta raw_tool_calls: Final = chunk.get("tool_calls", None) - tool_calls: Final = ( - [chunk["tool_use"]] if raw_tool_calls is None and chunk.get("tool_use") is not None else raw_tool_calls - ) + legacy_tool_use: Final = chunk.get("tool_use", None) + tool_calls: Final = [legacy_tool_use] if raw_tool_calls is None and legacy_tool_use is not None else raw_tool_calls raw_content: Final = chunk["text"] content: Final = ( diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index af9e7bc3c50..f9757eef4f8 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -54,8 +54,8 @@ class ModelResponseIterator: tool_calls_list: list[ChatCompletionToolCallChunk] = [] # mutable-ok: accumulates delta calls if choice_delta.tool_calls is not None: for tc in choice_delta.tool_calls: - tc_fn: Final = getattr(tc, "function", None) - fn_chunk: Final = ( + tc_fn = getattr(tc, "function", None) + fn_chunk = ( ChatCompletionToolCallFunctionChunk( name=getattr(tc_fn, "name", None), arguments=getattr(tc_fn, "arguments", "") or "", @@ -63,7 +63,7 @@ class ModelResponseIterator: if tc_fn is not None else ChatCompletionToolCallFunctionChunk(name=None, arguments="") ) - call_chunk: Final = ChatCompletionToolCallChunk( + call_chunk = ChatCompletionToolCallChunk( id=getattr(tc, "id", None), type="function", function=fn_chunk, diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index c43f91edf9d..d85f9bdb38f 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -222,3 +222,386 @@ def test_real_stream_keeps_reasoning_multi_tool_calls_and_terminal_usage(): assert chunks[4].usage.completion_tokens == 64 assert chunks[4].usage.prompt_tokens == 10 assert chunks[4].usage.completion_tokens_details.reasoning_tokens == 30 + + +def test_real_stream_incremental_alternating_tool_calls_and_negative_controls(): + events: Final[list[str]] = [ + # Frame 1: index 0 initial (id, name, arg prefix) + "data: " + + json.dumps( + { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "'}, + } + ] + }, + "finish_reason": None, + } + ], + } + ), + # Frame 2: index 1 initial (id, name, arg prefix) + "data: " + + json.dumps( + { + "id": "c2", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 1, + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"location": "'}, + } + ] + }, + "finish_reason": None, + } + ], + } + ), + # Frame 3: index 0 suffix (id and name omitted) + "data: " + + json.dumps( + { + "id": "c3", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "function": {"arguments": 'Paris"}'}, + } + ] + }, + "finish_reason": None, + } + ], + } + ), + # Frame 4: index 1 suffix (id and name omitted) + "data: " + + json.dumps( + { + "id": "c4", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 1, + "function": {"arguments": 'Europe/Paris"}'}, + } + ] + }, + "finish_reason": None, + } + ], + } + ), + # Frame 5: normal finish_reason length as reporter observed on Vertex + "data: " + + json.dumps( + { + "id": "c5", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [{"index": 0, "delta": {}, "finish_reason": "length"}], + } + ), + # Frame 6: empty-choice usage-only terminal completion_tokens=64 + "data: " + + json.dumps( + { + "id": "c6", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [], + "usage": { + "prompt_tokens": 12, + "completion_tokens": 64, + "total_tokens": 76, + "completion_tokens_details": {"reasoning_tokens": 18}, + }, + } + ), + "data: [DONE]", + ] + + sse_bytes: Final[bytes] = ("\n\n".join(events) + "\n\n").encode("utf-8") + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_bytes) + + mock_logging: Final = Mock() + mock_logging.model_call_details = {"litellm_params": {}} + mock_logging.stream_options = {"include_usage": True} + + raw_stream: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + api_base="https://openai-like.example/v1/chat/completions", + headers={}, + data=json.dumps({"model": "glm-4", "messages": []}), + model="glm-4", + messages=[], + logging_obj=mock_logging, + fake_stream=False, + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=raw_stream, + model="glm-4", + custom_llm_provider="vertex_ai", + logging_obj=mock_logging, + stream_options={"include_usage": True}, + ) + + chunks: Final = list(wrapper) + assert len(chunks) == 6 + + # Assert per-frame index/id/argument preservation + # Frame 0: index 0 initial + assert chunks[0].choices[0].delta.tool_calls[0].index == 0 + assert chunks[0].choices[0].delta.tool_calls[0].id == "call_weather" + assert chunks[0].choices[0].delta.tool_calls[0].function.name == "get_weather" + assert chunks[0].choices[0].delta.tool_calls[0].function.arguments == '{"city": "' + + # Frame 1: index 1 initial + assert chunks[1].choices[0].delta.tool_calls[0].index == 1 + assert chunks[1].choices[0].delta.tool_calls[0].id == "call_time" + assert chunks[1].choices[0].delta.tool_calls[0].function.name == "get_time" + assert chunks[1].choices[0].delta.tool_calls[0].function.arguments == '{"location": "' + + # Frame 2: index 0 continuation (id/name absent) + assert chunks[2].choices[0].delta.tool_calls[0].index == 0 + assert chunks[2].choices[0].delta.tool_calls[0].id is None + assert chunks[2].choices[0].delta.tool_calls[0].function.name is None + assert chunks[2].choices[0].delta.tool_calls[0].function.arguments == 'Paris"}' + + # Frame 3: index 1 continuation (id/name absent) + assert chunks[3].choices[0].delta.tool_calls[0].index == 1 + assert chunks[3].choices[0].delta.tool_calls[0].id is None + assert chunks[3].choices[0].delta.tool_calls[0].function.name is None + assert chunks[3].choices[0].delta.tool_calls[0].function.arguments == 'Europe/Paris"}' + + # Frame 4: finish_reason length + assert chunks[4].choices[0].finish_reason == "length" + + # Frame 5: usage chunk + assert chunks[5].usage is not None + assert chunks[5].usage.completion_tokens == 64 + assert chunks[5].usage.prompt_tokens == 12 + + # Reassembling in client: final separately ordered complete arguments + assembled_calls: Final[dict[int, dict[str, str]]] = {} + for c in chunks: + if c.choices and c.choices[0].delta and c.choices[0].delta.tool_calls: + for tc in c.choices[0].delta.tool_calls: + idx = tc.index + if idx not in assembled_calls: + assembled_calls[idx] = {"id": tc.id or "", "name": tc.function.name or "", "arguments": ""} + if tc.id: + assembled_calls[idx]["id"] = tc.id + if tc.function.name: + assembled_calls[idx]["name"] = tc.function.name + assembled_calls[idx]["arguments"] += tc.function.arguments + + assert assembled_calls[0]["id"] == "call_weather" + assert assembled_calls[0]["name"] == "get_weather" + assert assembled_calls[0]["arguments"] == '{"city": "Paris"}' + + assert assembled_calls[1]["id"] == "call_time" + assert assembled_calls[1]["name"] == "get_time" + assert assembled_calls[1]["arguments"] == '{"location": "Europe/Paris"}' + + +def test_real_stream_content_only_and_no_provider_usage_negative_controls(): + # Negative control 1: content-only stream emits no ghost tool calls, reasoning, or provider usage + events_content: Final[list[str]] = [ + "data: " + + json.dumps( + { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": {"content": "Direct text reply."}, + "finish_reason": None, + } + ], + } + ), + "data: " + + json.dumps( + { + "id": "c2", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + ), + "data: [DONE]", + ] + sse_bytes_content: Final[bytes] = ("\n\n".join(events_content) + "\n\n").encode("utf-8") + + def respond_content(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_bytes_content) + + mock_logging_content: Final = Mock() + mock_logging_content.model_call_details = {"litellm_params": {}} + mock_logging_content.stream_options = {} + + stream_content: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond_content))), + api_base="https://openai-like.example/v1/chat/completions", + headers={}, + data=json.dumps({"model": "glm-4", "messages": []}), + model="glm-4", + messages=[], + logging_obj=mock_logging_content, + fake_stream=False, + ) + wrapper_content: Final = CustomStreamWrapper( + completion_stream=stream_content, + model="glm-4", + custom_llm_provider="vertex_ai", + logging_obj=mock_logging_content, + stream_options={}, + ) + chunks_content: Final = list(wrapper_content) + assert len(chunks_content) == 2 + assert chunks_content[0].choices[0].delta.content == "Direct text reply." + assert chunks_content[0].choices[0].delta.tool_calls is None + assert getattr(chunks_content[0].choices[0].delta, "reasoning_content", None) is None + assert getattr(chunks_content[0], "usage", None) is None + assert chunks_content[1].choices[0].finish_reason == "stop" + + # Negative control 2: no-provider-usage stream terminates cleanly without synthetic usage chunks + events_no_usage: Final[list[str]] = [ + "data: " + + json.dumps( + { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_single", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + }, + "finish_reason": None, + } + ], + } + ), + "data: " + + json.dumps( + { + "id": "c2", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-4", + "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}], + } + ), + "data: [DONE]", + ] + sse_bytes_no_usage: Final[bytes] = ("\n\n".join(events_no_usage) + "\n\n").encode("utf-8") + + def respond_no_usage(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_bytes_no_usage) + + mock_logging_no_usage: Final = Mock() + mock_logging_no_usage.model_call_details = {"litellm_params": {}} + mock_logging_no_usage.stream_options = {} + + stream_no_usage: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond_no_usage))), + api_base="https://openai-like.example/v1/chat/completions", + headers={}, + data=json.dumps({"model": "glm-4", "messages": []}), + model="glm-4", + messages=[], + logging_obj=mock_logging_no_usage, + fake_stream=False, + ) + wrapper_no_usage: Final = CustomStreamWrapper( + completion_stream=stream_no_usage, + model="glm-4", + custom_llm_provider="vertex_ai", + logging_obj=mock_logging_no_usage, + stream_options={}, + ) + chunks_no_usage: Final = list(wrapper_no_usage) + assert len(chunks_no_usage) == 2 + assert chunks_no_usage[0].choices[0].delta.tool_calls[0].id == "call_single" + assert getattr(chunks_no_usage[0], "usage", None) is None + assert getattr(chunks_no_usage[1], "usage", None) is None + assert chunks_no_usage[1].choices[0].finish_reason == "tool_calls" + + +def test_legacy_tool_use_fallback_covered(): + from litellm.litellm_core_utils.streaming_handler import convert_generic_chunk_to_model_response_stream + from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ) + from litellm.types.utils import GenericStreamingChunk + + legacy_chunk: Final[GenericStreamingChunk] = GenericStreamingChunk( + text="", + tool_use=ChatCompletionToolCallChunk( + id="call_legacy", + type="function", + function=ChatCompletionToolCallFunctionChunk(name="old_tool", arguments="{}"), + index=0, + ), + is_finished=False, + finish_reason="", + usage=None, + index=0, + ) + res_stream: Final = convert_generic_chunk_to_model_response_stream(legacy_chunk) + assert res_stream.choices[0].delta.tool_calls is not None + assert res_stream.choices[0].delta.tool_calls[0].id == "call_legacy"