diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index aea136b913c..93304e0f356 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -454,10 +454,16 @@ 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 def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: return self.chunk_parser(json.loads(str_line)) + def _could_be_function_call(self, buffered: str) -> bool: + normalized: Final = "".join(buffered.split()) + prefix: Final = '{"name"' + return normalized.startswith(prefix) or prefix.startswith(normalized) + def _parse_buffered_function_call(self) -> ChatCompletionDeltaToolCall | None: if self.buffered_json_content is None: return None @@ -472,7 +478,9 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): type="function", function=Function( name=parsed["name"], - arguments=json.dumps(parsed["arguments"]), + arguments=( + parsed["arguments"] if isinstance(parsed["arguments"], str) else json.dumps(parsed["arguments"]) + ), ), ) return None @@ -531,15 +539,24 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): elif chunk["response"]: text = chunk["response"] if self.buffered_json_content is not None or ( - self.buffered_json_content is None + not self.function_call_buffering_disabled and not self.started_reasoning_content and text.lstrip().startswith("{") ): - self.buffered_json_content = (self.buffered_json_content or "") + text - return ModelResponseStream( - choices=[StreamingChoices(index=0, delta=Delta())], # mutable-ok: requires a list of choices - usage=None, - ) + 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 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 aa5fbd0faf5..6209f246d1d 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -540,13 +540,30 @@ class TestOllamaTextCompletionStreamingToolCalls: assert json.loads(tool_calls[0].function.arguments) == {"location": "Paris"} assert done.choices[0].finish_reason == "tool_calls" - def test_streamed_regular_json_emitted_as_content_on_done(self): + def test_streamed_regular_json_flushed_as_content_once_not_a_function_call(self): chunks, done = self._stream(['{"answer":', ' 42}']) + assert isinstance(chunks[0], ModelResponseStream) + assert chunks[0].choices[0].delta.content == '{"answer":' + assert chunks[1].choices[0].delta.content == " 42}" + assert done["finish_reason"] == "stop" + + def test_streamed_function_call_with_string_arguments_not_double_encoded(self): + chunks, done = self._stream(['{"name": "get_weather",', ' "arguments": "{\\"location\\": \\"Paris\\"}"}']) + assert isinstance(done, ModelResponseStream) - assert done.choices[0].delta.tool_calls is None - assert done.choices[0].delta.content == '{"answer": 42}' - assert done.choices[0].finish_reason == "stop" + tool_calls = done.choices[0].delta.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 + assert json.loads(tool_calls[0].function.arguments) == {"location": "Paris"} + + def test_brace_prefixed_prose_streams_incrementally(self): + chunks, done = self._stream(["{note: this", " is not JSON}", " and more text"]) + + assert isinstance(chunks[0], ModelResponseStream) + assert chunks[0].choices[0].delta.content == "{note: this" + assert chunks[1].choices[0].delta.content == " is not JSON}" + assert chunks[2].choices[0].delta.content == " and more text" + assert done["finish_reason"] == "stop" def test_plain_text_still_streams_incrementally(self): chunks, done = self._stream(["Hello", " world"])