From db1cc736addb39db761ae847810d85145a859c50 Mon Sep 17 00:00:00 2001 From: Arthi Arumugam Date: Mon, 27 Jul 2026 16:41:12 +0530 Subject: [PATCH] fix(ollama): emit streamed tool calls instead of raw JSON content Tools are not a supported OpenAI param on the ollama/ route, so get_optional_params rewrites them into a JSON-only prompt and sets format=json. transform_response closes that loop on the non-streaming path: it parses the finished body and turns a name/arguments object into a real tool_calls message with finish_reason tool_calls. OllamaTextCompletionResponseIterator.chunk_parser never had the second half, so streaming sent the JSON object out as content deltas and ended the turn with stop. The tool was never invoked and the raw JSON was rendered to the end user. chunk_parser only ever sees fragments, so it cannot parse mid-stream. Hold response text back while the accumulated text is still consistent with a JSON object, release everything held the moment it cannot be one, and decide at done. Prose fails that check on its first fragment, so plain text streaming is unaffected. By the time the request reaches the provider the tools survive only as prompt text, so get_optional_params passes on the functions it rewrote. transform_request consumes that, so it never reaches ollama, and it is what arms the iterator. A request that offered no tools never buffers a byte, json_object mode included, and a body naming a function the request did not offer stays content. The object itself must match what litellm asks ollama for in function_call_prompt: name, arguments, and nothing else. --- .../llms/ollama/completion/transformation.py | 105 +++++- litellm/utils.py | 1 + .../test_ollama_completion_transformation.py | 317 ++++++++++++++++++ 3 files changed, 419 insertions(+), 4 deletions(-) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 204b0d15c03..72089574b3b 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, Union from httpx._models import Headers, Response +from pydantic import BaseModel, ConfigDict, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -17,7 +18,12 @@ ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException -from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ChatCompletionUsageBlock, +) from litellm.types.utils import ( Delta, GenericStreamingChunk, @@ -99,6 +105,8 @@ class OllamaConfig(BaseConfig): system: Optional[str] = None template: Optional[str] = None + _prompted_tool_names: frozenset[str] = frozenset() + def __init__( self, mirostat: Optional[int] = None, @@ -376,6 +384,7 @@ def transform_request( format = optional_params.pop("format", None) images = optional_params.pop("images", None) think = optional_params.pop("think", None) + self._prompted_tool_names = _prompted_tool_names(optional_params.pop("prompted_tool_calls", None)) data = { "model": model, "prompt": ollama_prompt, @@ -439,14 +448,89 @@ def get_model_response_iterator( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, + tool_names=self._prompted_tool_names, ) +def _prompted_tool_names(prompted_tools: object) -> frozenset[str]: + """Names of the functions `get_optional_params` rewrote into this request's prompt.""" + if not isinstance(prompted_tools, list): + return frozenset() + + schemas = ((tool.get("function") or tool) for tool in prompted_tools if isinstance(tool, dict)) + return frozenset( + schema["name"] for schema in schemas if isinstance(schema, dict) and isinstance(schema.get("name"), str) + ) + + +class _OllamaJsonModeToolCall(BaseModel): + """ + The exact object litellm's own tool prompt asks ollama for, in `factory.function_call_prompt`. + Extra keys are rejected so a JSON body that merely happens to carry a `name` stays content. + """ + + model_config = ConfigDict(extra="forbid") + + name: str + arguments: dict[str, object] + + +def _build_tool_call(response_text: str, tool_names: frozenset[str]) -> ChatCompletionToolCallChunk | None: + """ + Build a tool call out of a completed response body, but only for a function this request + actually offered. This is a stricter test than the one `OllamaConfig.transform_response` + applies to the same body, so it never converts anything that path would leave as content. + """ + try: + function_call = _OllamaJsonModeToolCall.model_validate_json(response_text) + except ValidationError: + return None + + if function_call.name not in tool_names: + return None + + return ChatCompletionToolCallChunk( + id=f"call_{uuid.uuid4()}", + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=function_call.name, + arguments=json.dumps(function_call.arguments), + ), + index=0, + ) + + class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): - def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): + def __init__( + self, + streaming_response, + sync_stream: bool, + json_mode: Optional[bool] = False, + tool_names: frozenset[str] = frozenset(), + ): super().__init__(streaming_response, sync_stream, json_mode) self.started_reasoning_content: bool = False self.finished_reasoning_content: bool = False + self.tool_names: frozenset[str] = tool_names + self.tool_call_buffer: str | None = "" if tool_names else None + + def _hold_back_or_release(self, content: str) -> str | None: + """ + `/api/generate` returns a tool call as a JSON object inside the response text, so it can + only be recognised once that object is complete. Hold text back while it can still turn + out to be one, and release everything held as soon as it cannot. + """ + if self.tool_call_buffer is None: + return content + + self.tool_call_buffer += content + stripped = self.tool_call_buffer.lstrip() + if not stripped or stripped.startswith("{"): + return None + + released = self.tool_call_buffer + self.tool_call_buffer = None + return released def _handle_string_chunk(self, str_line: str) -> Union[GenericStreamingChunk, ModelResponseStream]: return self.chunk_parser(json.loads(str_line)) @@ -473,8 +557,21 @@ def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelRespons completion_tokens=eval_count, total_tokens=prompt_eval_count + eval_count, ) + + held_back = self.tool_call_buffer + self.tool_call_buffer = None + tool_call = _build_tool_call(held_back, self.tool_names) if held_back else None + if tool_call is not None: + return GenericStreamingChunk( + text="", + tool_use=tool_call, + is_finished=True, + finish_reason="tool_calls", + usage=usage, + ) + return GenericStreamingChunk( - text=text, + text=held_back or text, is_finished=is_finished, finish_reason=finish_reason, usage=usage, @@ -494,7 +591,7 @@ def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelRespons if self.started_reasoning_content and not self.finished_reasoning_content: reasoning_content = text else: - content = text + content = self._hold_back_or_release(text) return ModelResponseStream( choices=[ diff --git a/litellm/utils.py b/litellm/utils.py index a11c5500503..3403a40168b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3746,6 +3746,7 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c non_default_params.pop("tool_choice", None) # causes ollama requests to hang elif "functions" in non_default_params: optional_params["functions_unsupported_model"] = non_default_params.pop("functions") + optional_params["prompted_tool_calls"] = optional_params.get("functions_unsupported_model") or [] elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead optional_params["functions_unsupported_model"] = non_default_params.pop( "tools", non_default_params.pop("functions", 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 dd59cdcac1c..38edb1a704b 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -507,3 +507,320 @@ def test_chunk_parser_done_chunk(self): assert result["usage"]["prompt_tokens"] == 10 assert result["usage"]["completion_tokens"] == 5 assert result["usage"]["total_tokens"] == 15 + + def test_chunk_parser_streams_tool_call_instead_of_raw_json(self): + """A format=json tool call must arrive as tool_calls, not as JSON in the content.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + tool_names=frozenset({"lookup_account", "lookup"}), + ) + + fragments = [ + '{"', + "name", + '": "', + "lookup_account", + '", "', + "arguments", + '": {"', + "email", + '": "', + "maya.iyer@example.com", + '"}}', + ] + streamed_content = "" + for fragment in fragments: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + assert streamed_content == "" + + result = iterator.chunk_parser( + { + "model": "qwen2.5:3b", + "response": "", + "done": True, + "prompt_eval_count": 120, + "eval_count": 24, + } + ) + + assert result["text"] == "" + assert result["finish_reason"] == "tool_calls" + assert result["is_finished"] is True + assert result["tool_use"] is not None + assert result["tool_use"]["type"] == "function" + assert result["tool_use"]["index"] == 0 + assert result["tool_use"]["id"] + assert result["tool_use"]["function"]["name"] == "lookup_account" + assert json.loads(result["tool_use"]["function"]["arguments"]) == {"email": "maya.iyer@example.com"} + + def test_chunk_parser_streams_prose_incrementally(self): + """Prose must keep streaming fragment by fragment and carry no tool call.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + tool_names=frozenset({"lookup_account", "lookup"}), + ) + + fragments = ["I ", "will ", "look ", "that ", "up", "."] + deltas = [] + for fragment in fragments: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + deltas.append(result.choices[0].delta.content) + + assert deltas == fragments + + result = iterator.chunk_parser( + { + "model": "qwen2.5:3b", + "response": "", + "done": True, + "prompt_eval_count": 10, + "eval_count": 6, + } + ) + + assert result["text"] == "" + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + def test_chunk_parser_releases_non_tool_call_json_as_content(self): + """A JSON body that is not a tool call must still be delivered as content.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + tool_names=frozenset({"lookup_account", "lookup"}), + ) + + fragments = ['{"', "city", '": "', "Chennai", '"}'] + streamed_content = "" + for fragment in fragments: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + result = iterator.chunk_parser( + { + "model": "qwen2.5:3b", + "response": "", + "done": True, + "prompt_eval_count": 10, + "eval_count": 6, + } + ) + + assert streamed_content + result["text"] == '{"city": "Chennai"}' + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + def test_chunk_parser_leaves_json_carrying_extra_keys_as_content(self): + """A JSON body that merely happens to carry a name must not be mistaken for a tool call.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + tool_names=frozenset({"lookup_account", "lookup"}), + ) + + body = '{"name": "Chennai", "arguments": {"x": 1}, "population": 7000000}' + streamed_content = "" + for fragment in [body[i : i + 6] for i in range(0, len(body), 6)]: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + result = iterator.chunk_parser( + {"model": "qwen2.5:3b", "response": "", "done": True, "prompt_eval_count": 5, "eval_count": 9} + ) + + assert streamed_content + result["text"] == body + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + def test_chunk_parser_leaves_non_object_arguments_as_content(self): + """`arguments` must be the object the tool prompt asks for, not any value that happens to be there.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + tool_names=frozenset({"lookup_account", "lookup"}), + ) + + body = '{"name": "lookup", "arguments": "not an object"}' + streamed_content = "" + for fragment in [body[i : i + 6] for i in range(0, len(body), 6)]: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + result = iterator.chunk_parser( + {"model": "qwen2.5:3b", "response": "", "done": True, "prompt_eval_count": 5, "eval_count": 9} + ) + + assert streamed_content + result["text"] == body + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + +TOOL_CALL_BODY = '{"name": "lookup_account", "arguments": {"email": "maya.iyer@example.com"}}' + +LOOKUP_ACCOUNT_TOOL = { + "type": "function", + "function": { + "name": "lookup_account", + "description": "Look up a customer account by email address", + "parameters": {"type": "object", "properties": {"email": {"type": "string"}}, "required": ["email"]}, + }, +} + + +def _generate_transport(): + """Stands in for ollama's /api/generate, fragmenting the body the way it really does.""" + import httpx + + done = {"model": "qwen2.5:3b", "response": "", "done": True, "prompt_eval_count": 146, "eval_count": 21} + + def handler(request: "httpx.Request") -> "httpx.Response": + fragments = [TOOL_CALL_BODY[i : i + 4] for i in range(0, len(TOOL_CALL_BODY), 4)] + lines = [json.dumps({"model": "qwen2.5:3b", "response": f, "done": False}) for f in fragments] + lines.append(json.dumps(done)) + return httpx.Response(200, content=("\n".join(lines) + "\n").encode()) + + return httpx.MockTransport(handler) + + +class TestOllamaStreamingToolCallsEndToEnd: + def test_streamed_tool_call_is_not_delivered_as_content(self): + """The reported bug: a streamed tool call reaching the caller as raw JSON in the content.""" + import httpx + + from litellm import completion + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler(client=httpx.Client(transport=_generate_transport())) + + response = completion( + model="ollama/qwen2.5:3b", + messages=[{"role": "user", "content": "Cancel the subscription for maya.iyer@example.com"}], + tools=[LOOKUP_ACCOUNT_TOOL], + api_base="http://localhost:11434", + stream=True, + client=client, + ) + + streamed_content = "" + tool_calls = [] + finish_reason = None + for chunk in response: + delta = chunk.choices[0].delta + streamed_content += getattr(delta, "content", None) or "" + tool_calls += getattr(delta, "tool_calls", None) or [] + finish_reason = chunk.choices[0].finish_reason or finish_reason + + assert streamed_content == "" + assert finish_reason == "tool_calls" + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "lookup_account" + assert json.loads(tool_calls[0].function.arguments) == {"email": "maya.iyer@example.com"} + + def test_detection_is_off_for_a_request_that_offered_no_tools(self): + """No tools were rewritten into a prompt, so a tool-shaped body is just content.""" + iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) + + streamed_content = "" + for fragment in [TOOL_CALL_BODY[i : i + 6] for i in range(0, len(TOOL_CALL_BODY), 6)]: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + result = iterator.chunk_parser( + {"model": "qwen2.5:3b", "response": "", "done": True, "prompt_eval_count": 5, "eval_count": 9} + ) + + assert streamed_content == TOOL_CALL_BODY + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + def test_detection_is_armed_only_with_the_functions_the_request_offered(self): + """`get_optional_params` passes on the functions it rewrote; those are the only ones accepted.""" + from litellm.utils import get_optional_params + + config = OllamaConfig() + base = { + "model": "qwen2.5:3b", + "messages": [{"role": "user", "content": "hi"}], + "litellm_params": {}, + "headers": {}, + } + + json_mode_only = get_optional_params( + model="qwen2.5:3b", + custom_llm_provider="ollama", + stream=True, + response_format={"type": "json_object"}, + ) + assert "prompted_tool_calls" not in json_mode_only + config.transform_request(optional_params=json_mode_only, **base) + assert config.get_model_response_iterator(iter([]), sync_stream=True).tool_names == frozenset() + + with_tools = get_optional_params( + model="qwen2.5:3b", + custom_llm_provider="ollama", + stream=True, + tools=[LOOKUP_ACCOUNT_TOOL], + ) + assert with_tools["prompted_tool_calls"] == [LOOKUP_ACCOUNT_TOOL] + config.transform_request(optional_params=with_tools, **base) + assert config.get_model_response_iterator(iter([]), sync_stream=True).tool_names == frozenset( + {"lookup_account"} + ) + + def test_a_function_the_request_never_offered_stays_content(self): + """A body naming some other function must not be synthesised into a tool call.""" + iterator = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), sync_stream=True, json_mode=False, tool_names=frozenset({"lookup_account"}) + ) + + body = '{"name": "delete_account", "arguments": {"email": "maya.iyer@example.com"}}' + streamed_content = "" + for fragment in [body[i : i + 6] for i in range(0, len(body), 6)]: + result = iterator.chunk_parser({"model": "qwen2.5:3b", "response": fragment, "done": False}) + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + streamed_content += result.choices[0].delta.content or "" + + result = iterator.chunk_parser( + {"model": "qwen2.5:3b", "response": "", "done": True, "prompt_eval_count": 5, "eval_count": 9} + ) + + assert streamed_content + result["text"] == body + assert result["finish_reason"] == "stop" + assert result.get("tool_use") is None + + def test_prompted_tool_calls_marker_is_not_sent_to_ollama(self): + """The marker is litellm-internal; it must not leak into the request body.""" + config = OllamaConfig() + + data = config.transform_request( + model="qwen2.5:3b", + messages=[{"role": "user", "content": "hi"}], + optional_params={"stream": True, "format": "json", "prompted_tool_calls": True}, + litellm_params={}, + headers={}, + ) + + assert "prompted_tool_calls" not in data + assert "prompted_tool_calls" not in data["options"]