From d32e9748c06e43a221da90aca85d5afe4881649c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 2 Sep 2026 18:43:00 -0400 Subject: [PATCH] fix(ollama): only reconstruct streaming tool calls when JSON output was requested The /api/generate streaming buffer reconstructed a prompted function call from any response that started with {"name", with no check that the caller had asked for one. A plain stream with no tools and no response_format whose output happened to parse as {"name": ..., "arguments": ...} lost its content entirely and came back as a synthesized tool call with finish_reason=tool_calls. The non-streaming path already gates that reconstruction on format=json, which litellm sets for ollama whenever tools are passed. Carry the same signal from the request transform into the streaming iterator so both paths agree. The flag rides on the config rather than the iterator's json_mode argument because the async streaming handler builds the iterator without forwarding json_mode, so a request through the proxy would never see it. --- .../llms/ollama/completion/transformation.py | 51 ++++++++------- .../test_ollama_completion_transformation.py | 63 ++++++++++++++++++- 2 files changed, 90 insertions(+), 24 deletions(-) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 93304e0f356..0eb7c66fa26 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -103,6 +103,7 @@ class OllamaConfig(BaseConfig): top_p: float | None = None system: str | None = None template: str | None = None + json_output_requested: bool = False def __init__( self, @@ -380,6 +381,7 @@ class OllamaConfig(BaseConfig): ollama_prompt = modified_prompt stream: Final = optional_params.pop("stream", False) format: Final = optional_params.pop("format", None) + self.json_output_requested = format == "json" images = optional_params.pop("images", None) think: Final = optional_params.pop("think", None) data: Final = { @@ -444,7 +446,7 @@ class OllamaConfig(BaseConfig): return OllamaTextCompletionResponseIterator( streaming_response=streaming_response, sync_stream=sync_stream, - json_mode=json_mode, + json_mode=json_mode or self.json_output_requested, ) @@ -454,7 +456,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): self.started_reasoning_content: bool = False self.finished_reasoning_content: bool = False self.buffered_json_content: str | None = None - self.function_call_buffering_disabled: bool = False + self.function_call_buffering_enabled: bool = bool(json_mode) def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: return self.chunk_parser(json.loads(str_line)) @@ -464,6 +466,23 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): prefix: Final = '{"name"' return normalized.startswith(prefix) or prefix.startswith(normalized) + def _released_text(self, response_text: str) -> str | None: + """None while a fragment is held back because it may still complete a prompted function call.""" + if self.buffered_json_content is None and not ( + self.function_call_buffering_enabled + and not self.started_reasoning_content + and response_text.lstrip().startswith("{") + ): + self.function_call_buffering_enabled = False + return response_text + candidate: Final = (self.buffered_json_content or "") + response_text + if self._could_be_function_call(candidate): + self.buffered_json_content = candidate + return None + self.buffered_json_content = None + self.function_call_buffering_enabled = False + return candidate + def _parse_buffered_function_call(self) -> ChatCompletionDeltaToolCall | None: if self.buffered_json_content is None: return None @@ -537,26 +556,14 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): usage=usage, ) elif chunk["response"]: - text = chunk["response"] - if self.buffered_json_content is not None or ( - not self.function_call_buffering_disabled - and not self.started_reasoning_content - and text.lstrip().startswith("{") - ): - candidate: Final = (self.buffered_json_content or "") + text - if self._could_be_function_call(candidate): - self.buffered_json_content = candidate - return ModelResponseStream( - choices=[ # mutable-ok: ModelResponseStream only accepts a list of choices - StreamingChoices(index=0, delta=Delta()) - ], - usage=None, - ) - self.buffered_json_content = None - self.function_call_buffering_disabled = True - text = candidate - else: - self.function_call_buffering_disabled = True + text = self._released_text(chunk["response"]) + if text is None: + return ModelResponseStream( + choices=[ # mutable-ok: ModelResponseStream only accepts a list of choices + StreamingChoices(index=0, delta=Delta()) + ], + usage=None, + ) reasoning_content: str | None = None content: str | None = None if text is not None: diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index 6209f246d1d..5ee3842cae5 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -507,8 +507,10 @@ class TestOllamaTextCompletionResponseIterator: class TestOllamaTextCompletionStreamingToolCalls: """Regression tests for https://github.com/BerriAI/litellm/issues/35711""" - def _stream(self, responses): - iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + def _stream(self, responses, json_mode=True): + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), sync_stream=True, json_mode=json_mode + ) chunks = [ iterator.chunk_parser({"model": "qwen3", "created_at": "t", "done": False, "response": r}) for r in responses @@ -571,3 +573,60 @@ class TestOllamaTextCompletionStreamingToolCalls: assert chunks[0].choices[0].delta.content == "Hello" assert chunks[1].choices[0].delta.content == " world" assert done["finish_reason"] == "stop" + + @pytest.mark.parametrize( + "arguments_fragment", + [' "arguments": {"location": "Paris"}}', ' "arguments": "{\\"location\\": \\"Paris\\"}"}'], + ) + def test_no_tool_call_reconstruction_when_json_was_not_requested(self, arguments_fragment): + """A caller that sent no tools and no response_format must never get a synthesized tool call, + and must never lose the content it did ask for.""" + chunks, done = self._stream(['{"name": "get_weather",', arguments_fragment], json_mode=False) + + streamed = "".join(c.choices[0].delta.content or "" for c in chunks) + assert streamed == '{"name": "get_weather",' + arguments_fragment + for chunk in chunks: + assert chunk.choices[0].delta.tool_calls is None + assert done["finish_reason"] == "stop" + + +class TestOllamaStreamGating: + """`utils.py` sets format=json for ollama whenever tools are passed, and the prompted function call + only makes sense for those requests. The iterator learns about it from the request transform.""" + + def _iterator_for(self, optional_params): + config = OllamaConfig() + config.transform_request( + model="qwen3", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + return config.get_model_response_iterator(streaming_response=iter([]), sync_stream=True, json_mode=False) + + def test_json_format_request_buffers_a_possible_function_call(self): + iterator = self._iterator_for({"format": "json"}) + + assert iterator.function_call_buffering_enabled is True + + def test_plain_request_never_buffers(self): + iterator = self._iterator_for({"temperature": 0.5}) + + assert iterator.function_call_buffering_enabled is False + + @pytest.mark.parametrize("sync_stream", [True, False]) + def test_gate_survives_both_sync_and_async_streaming(self, sync_stream): + """The async handler builds the iterator without forwarding json_mode, so the flag has to ride + on the config rather than on that argument.""" + config = OllamaConfig() + config.transform_request( + model="qwen3", + messages=[{"role": "user", "content": "hi"}], + optional_params={"format": "json"}, + litellm_params={}, + headers={}, + ) + iterator = config.get_model_response_iterator(streaming_response=iter([]), sync_stream=sync_stream) + + assert iterator.function_call_buffering_enabled is True