diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index e612faaa97d..ba59e3fa997 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5126,15 +5126,9 @@ def _bedrock_tools_pt(tools: list, model: str | None = None) -> list[BedrockTool return tool_block_list -FUNCTION_CALL_PROMPT_PREFIX: Final = ( - 'Produce JSON OUTPUT ONLY! Adhere to this format {"name": "function_name", ' - '"arguments":{"argument_name": "argument_value"}} The following functions are available to you:' -) - - # Function call template def function_call_prompt(messages: list, functions: list): - function_prompt = FUNCTION_CALL_PROMPT_PREFIX + function_prompt = """Produce JSON OUTPUT ONLY! Adhere to this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}} The following functions are available to you:""" for function in functions: function_prompt += f"""\n{function}\n""" diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 39ae2411997..c89d2571c45 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -1,4 +1,5 @@ import json +import re import time from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final @@ -12,7 +13,6 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.litellm_core_utils.prompt_templates.factory import ( - FUNCTION_CALL_PROMPT_PREFIX, convert_to_ollama_image, custom_prompt, ollama_pt, @@ -382,7 +382,7 @@ class OllamaConfig(BaseConfig): ollama_prompt = modified_prompt stream: Final = optional_params.pop("stream", False) format: Final = optional_params.pop("format", None) - self.function_call_prompted = format == "json" and FUNCTION_CALL_PROMPT_PREFIX in ollama_prompt + self.function_call_prompted = optional_params.pop("function_call_prompted", False) is True images = optional_params.pop("images", None) think: Final = optional_params.pop("think", None) data: Final = { @@ -452,6 +452,11 @@ class OllamaConfig(BaseConfig): ) +_FUNCTION_CALL_OPENING: Final = '{"name":"' +_FUNCTION_CALL_NAME_HEAD: Final = re.compile(r'^\{"name":"(?:[^"\\]|\\.)*"') +_FUNCTION_CALL_ARGUMENTS_KEY: Final = ',"arguments":' + + class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): def __init__( self, @@ -471,8 +476,13 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): 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) + name_head: Final = _FUNCTION_CALL_NAME_HEAD.match(normalized) + if name_head is None: + return _FUNCTION_CALL_OPENING.startswith(normalized) or normalized.startswith(_FUNCTION_CALL_OPENING) + after_name: Final = normalized[name_head.end() :] + return _FUNCTION_CALL_ARGUMENTS_KEY.startswith(after_name) or after_name.startswith( + _FUNCTION_CALL_ARGUMENTS_KEY + ) 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.""" diff --git a/litellm/utils.py b/litellm/utils.py index ba456fc353b..0a6722f5f23 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4108,6 +4108,7 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c if custom_llm_provider == "ollama": # ollama actually supports json output optional_params["format"] = "json" + optional_params["function_call_prompted"] = True litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt if "tools" in non_default_params: optional_params["functions_unsupported_model"] = non_default_params.pop("tools") 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 110c4ce751f..7cb039493ef 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -10,6 +10,7 @@ from litellm.llms.ollama.completion.transformation import ( OllamaTextCompletionResponseIterator, ) from litellm.types.utils import Message, ModelResponse, ModelResponseStream +from litellm.utils import get_optional_params class TestOllamaConfig: @@ -69,9 +70,7 @@ class TestOllamaConfig: # Create mock response with JSON function call format raw_response = MagicMock() raw_response.json.return_value = { - "response": json.dumps( - {"name": "get_weather", "arguments": {"location": "San Francisco"}} - ) + "response": json.dumps({"name": "get_weather", "arguments": {"location": "San Francisco"}}) } # Create properly structured model response object @@ -102,13 +101,10 @@ class TestOllamaConfig: assert result.choices[0]["finish_reason"] == "tool_calls" assert len(result.choices[0]["message"].tool_calls) == 1 assert result.choices[0]["message"].tool_calls[0]["id"].startswith("call_") - assert ( - result.choices[0]["message"].tool_calls[0]["function"]["name"] - == "get_weather" - ) - assert json.loads( - result.choices[0]["message"].tool_calls[0]["function"]["arguments"] - ) == {"location": "San Francisco"} + assert result.choices[0]["message"].tool_calls[0]["function"]["name"] == "get_weather" + assert json.loads(result.choices[0]["message"].tool_calls[0]["function"]["arguments"]) == { + "location": "San Francisco" + } # No usage assertions here as we don't need to test them in every case def test_transform_response_regular_json(self): @@ -118,9 +114,7 @@ class TestOllamaConfig: # Create mock response with regular JSON (not function call) raw_response = MagicMock() raw_response.json.return_value = { - "response": json.dumps( - {"result": "success", "data": {"temperature": 72, "unit": "F"}} - ) + "response": json.dumps({"result": "success", "data": {"temperature": 72, "unit": "F"}}) } # Create properly structured model response object @@ -147,9 +141,7 @@ class TestOllamaConfig: ) # Verify result has JSON content - expected_content = json.dumps( - {"result": "success", "data": {"temperature": 72, "unit": "F"}} - ) + expected_content = json.dumps({"result": "success", "data": {"temperature": 72, "unit": "F"}}) assert result.choices[0]["message"].content == expected_content assert result.choices[0]["finish_reason"] == "stop" # No usage assertions here as we don't need to test them in every case @@ -191,10 +183,7 @@ class TestOllamaConfig: ) # Verify reasoning content is extracted - assert ( - result.choices[0]["message"].reasoning_content - == "I need to think about this problem step by step" - ) + assert result.choices[0]["message"].reasoning_content == "I need to think about this problem step by step" assert result.choices[0]["message"].content == "Here is my answer" assert result.choices[0]["finish_reason"] == "stop" @@ -233,10 +222,7 @@ class TestOllamaConfig: ) # Verify reasoning content is extracted - assert ( - result.choices[0]["message"].reasoning_content - == "Let me analyze this carefully" - ) + assert result.choices[0]["message"].reasoning_content == "Let me analyze this carefully" assert result.choices[0]["message"].content == "The solution is X" assert result.choices[0]["finish_reason"] == "stop" @@ -277,10 +263,7 @@ class TestOllamaConfig: # Verify multiline reasoning content is extracted expected_reasoning = "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\n" assert result.choices[0]["message"].reasoning_content == expected_reasoning - assert ( - result.choices[0]["message"].content - == "Based on my analysis, the answer is Y" - ) + assert result.choices[0]["message"].content == "Based on my analysis, the answer is Y" assert result.choices[0]["finish_reason"] == "stop" def test_transform_response_thinking_only(self): @@ -318,10 +301,7 @@ class TestOllamaConfig: ) # Verify reasoning content is extracted and content is empty - assert ( - result.choices[0]["message"].reasoning_content - == "Just internal thoughts, no response" - ) + assert result.choices[0]["message"].reasoning_content == "Just internal thoughts, no response" assert result.choices[0]["message"].content == "" assert result.choices[0]["finish_reason"] == "stop" @@ -360,10 +340,7 @@ class TestOllamaConfig: ) # Verify reasoning content is extracted even in JSON mode when JSON parsing fails - assert ( - result.choices[0]["message"].reasoning_content - == "Planning my JSON response" - ) + assert result.choices[0]["message"].reasoning_content == "Planning my JSON response" assert result.choices[0]["message"].content == "This is not valid JSON" assert result.choices[0]["finish_reason"] == "stop" @@ -403,19 +380,14 @@ class TestOllamaConfig: # Verify no reasoning content is extracted assert result.choices[0]["message"].reasoning_content is None - assert ( - result.choices[0]["message"].content - == "Regular response without any thinking tags" - ) + assert result.choices[0]["message"].content == "Regular response without any thinking tags" assert result.choices[0]["finish_reason"] == "stop" class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_with_thinking_field(self): """Test that chunks with 'thinking' field and empty 'response' are handled correctly.""" - iterator = OllamaTextCompletionResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) # Test chunk with thinking field - this is the problematic case from the issue chunk_with_thinking = { @@ -435,9 +407,7 @@ class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_normal_response(self): """Test that normal response chunks still work.""" - iterator = OllamaTextCompletionResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) # Test normal chunk with response normal_chunk = { @@ -457,9 +427,7 @@ class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_empty_response_without_thinking(self): """Test that empty response chunks without thinking still work.""" - iterator = OllamaTextCompletionResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) # Test empty response chunk without thinking empty_response_chunk = { @@ -479,9 +447,7 @@ class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_done_chunk(self): """Test that done chunks work correctly.""" - iterator = OllamaTextCompletionResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) # Test done chunk done_chunk = { @@ -563,8 +529,19 @@ class TestOllamaTextCompletionStreamingToolCalls: assert streamed == "\nHello world" assert done["finish_reason"] == "stop" + def test_name_prefixed_regular_json_flushed_as_content_once_arguments_key_is_ruled_out(self): + """`{"name": ...}` alone is not enough to hold the stream: once the key after `name` is not + `arguments`, the buffered text must be delivered before the terminal chunk.""" + chunks, done = self._stream(['{"name": "Alice",', ' "age": 30}', " extra"]) + + assert not chunks[0].choices[0].delta.content + assert chunks[1].choices[0].delta.content == '{"name": "Alice", "age": 30}' + assert chunks[2].choices[0].delta.content == " extra" + assert all(c.choices[0].delta.tool_calls is None for c in chunks) + assert done["finish_reason"] == "stop" + def test_streamed_regular_json_flushed_as_content_once_not_a_function_call(self): - chunks, done = self._stream(['{"answer":', ' 42}']) + chunks, done = self._stream(['{"answer":', " 42}"]) assert isinstance(chunks[0], ModelResponseStream) assert chunks[0].choices[0].delta.content == '{"answer":' @@ -612,75 +589,81 @@ class TestOllamaTextCompletionStreamingToolCalls: class TestOllamaStreamGating: - """`utils.py` sets format=json and injects `function_call_prompt` into the messages whenever tools are - passed to `ollama/`. Only those requests may have their streamed JSON reconstructed into a tool call; - the iterator learns about it from the request transform.""" + """`get_optional_params` marks `ollama/` requests that carried tools, since it is the one turning them + into a prompted function call plus format=json. Only those requests may have their streamed JSON + reconstructed into a tool call; `response_format=json_object` also sets format=json but the caller + asked for JSON content there and must get it back as content.""" - _tool = {"name": "get_weather", "parameters": {"type": "object", "properties": {}}} + _tools = [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + } + ] + _function_shaped_json = ['{"name": "get_weather",', ' "arguments": {"location": "Paris"}}'] - def _iterator_for(self, optional_params, messages=None, sync_stream=True): + def _stream_through_config(self, optional_params, messages=None, sync_stream=True): config = OllamaConfig() - config.transform_request( + request = config.transform_request( model="qwen3", messages=messages if messages is not None else [{"role": "user", "content": "hi"}], optional_params=optional_params, litellm_params={}, headers={}, ) - return config.get_model_response_iterator(streaming_response=iter([]), sync_stream=sync_stream) - - def test_tools_request_buffers_a_possible_function_call(self): - messages = function_call_prompt([{"role": "user", "content": "weather in Paris?"}], [self._tool]) - - iterator = self._iterator_for({"format": "json"}, messages=messages) - - assert iterator.function_call_buffering_enabled is True - - def test_tools_request_with_existing_system_message_buffers(self): - messages = function_call_prompt( - [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "weather?"}], [self._tool] - ) - - iterator = self._iterator_for({"format": "json"}, messages=messages) - - assert iterator.function_call_buffering_enabled is True - - def test_json_output_without_tools_never_buffers(self): - """response_format=json_object alone must not turn `{"name": ..., "arguments": ...}` content into - a tool call.""" - iterator = self._iterator_for({"format": "json"}) - - assert iterator.function_call_buffering_enabled is False - - 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.""" - messages = function_call_prompt([{"role": "user", "content": "hi"}], [self._tool]) - - iterator = self._iterator_for({"format": "json"}, messages=messages, sync_stream=sync_stream) - - assert iterator.function_call_buffering_enabled is True - - def test_json_output_without_tools_streams_name_arguments_object_as_content(self): - """End to end through the config: a tool-free JSON request whose schema happens to use top-level - `name` and `arguments` keys keeps its content and gets no tool call.""" - iterator = self._iterator_for({"format": "json"}) - fragments = ['{"name": "Alice",', ' "arguments": ["x"]}'] - + iterator = config.get_model_response_iterator(streaming_response=iter([]), sync_stream=sync_stream) chunks = [ iterator.chunk_parser({"model": "qwen3", "created_at": "t", "done": False, "response": r}) - for r in fragments + for r in self._function_shaped_json ] done = iterator.chunk_parser( {"model": "qwen3", "created_at": "t", "done": True, "done_reason": "stop", "response": ""} ) + return request, chunks, done - assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(fragments) + def _assert_streamed_as_content(self, chunks, done): + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(self._function_shaped_json) assert all(c.choices[0].delta.tool_calls is None for c in chunks) assert done["finish_reason"] == "stop" + + @pytest.mark.parametrize("sync_stream", [True, False]) + def test_tools_request_reconstructs_the_streamed_function_call(self, sync_stream): + """The async handler builds the iterator without forwarding json_mode, so the gate has to ride + on the config rather than on that argument.""" + optional_params = get_optional_params(model="qwen3", custom_llm_provider="ollama", tools=self._tools) + + request, chunks, done = self._stream_through_config(optional_params, sync_stream=sync_stream) + + assert request["format"] == "json" + assert "function_call_prompted" not in request["options"] + assert all(not chunk.choices[0].delta.content for chunk in chunks) + tool_calls = done.choices[0].delta.tool_calls + assert tool_calls is not None and tool_calls[0].function.name == "get_weather" + assert done.choices[0].finish_reason == "tool_calls" + + def test_json_object_response_format_without_tools_keeps_function_shaped_json_as_content(self): + optional_params = get_optional_params( + model="qwen3", custom_llm_provider="ollama", response_format={"type": "json_object"} + ) + + request, chunks, done = self._stream_through_config(optional_params) + + assert request["format"] == "json" + self._assert_streamed_as_content(chunks, done) + + def test_function_call_prompt_text_in_messages_does_not_enable_reconstruction(self): + """Only the request carrying tools may buffer. A caller pasting litellm's own function-call prompt + into its messages while asking for JSON output still gets its JSON back as content.""" + optional_params = get_optional_params( + model="qwen3", custom_llm_provider="ollama", response_format={"type": "json_object"} + ) + messages = function_call_prompt([{"role": "user", "content": "weather in Paris?"}], [self._tools[0]]) + + _, chunks, done = self._stream_through_config(optional_params, messages=messages) + + self._assert_streamed_as_content(chunks, done) + + def test_plain_request_never_buffers(self): + _, chunks, done = self._stream_through_config({"temperature": 0.5}) + + self._assert_streamed_as_content(chunks, done)