scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(ollama): gate tool call reconstruction on the request marker and release name-prefixed JSON early

Browse source 
The prompt-prefix check could be satisfied by a caller pasting litellm's own
function-call prompt into its messages, so get_optional_params now marks the
request it converts into a prompted function call and transform_request reads
that marker instead of sniffing the prompt. The buffering heuristic also stops
matching once the key after "name" is anything other than "arguments", so
ordinary JSON such as {"name": "Alice", "age": 30} streams incrementally
again instead of being held until the terminal chunk
This commit is contained in:
mubashir1osmani 2026-09-02 21:00:31 -04:00
parent f516c01128
commit 5289bf2fcc
4 changed files with 106 additions and 118 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -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"""

18
litellm/llms/ollama/completion/transformation.py
View file

@ -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."""

1
litellm/utils.py
View file

@ -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")

197
tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py
View file

@ -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)