fix(ollama): address PR review feedback

litellm/llms/ollama/completion/transformation.py

@@ -1,6 +1,6 @@
 import json
 import time
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 from httpx._models import Headers, Response
@@ -189,10 +189,7 @@ class OllamaConfig(BaseConfig):
                     optional_params["format"] = value["json_schema"]["schema"]
 
         if "functions_unsupported_model" in optional_params and optional_params.get("stream") is True:
-            # Tools are emulated via a prompt instruction + format=json for ollama/, not a
-            # native `tools` request field. transform_response() already reconstructs the
-            # resulting JSON into tool_calls for stream=False; fake_stream reuses that same
-            # reconstruction for stream=True instead of forwarding the raw JSON as text.
+            # functions_unsupported_model means tools are emulated via prompt injection here.
             optional_params["fake_stream"] = True
 
         return optional_params
@@ -404,15 +401,15 @@ class OllamaConfig(BaseConfig):
 
     def sign_request(
         self,
-        headers: dict,  # mutable-ok: matches BaseConfig.sign_request's fixed override signature
-        optional_params: dict,  # mutable-ok: matches BaseConfig.sign_request's fixed override signature
-        request_data: dict,  # mutable-ok: matches BaseConfig.sign_request's fixed override signature
+        headers: dict[str, str],  # mutable-ok: returned unchanged, must stay assignable to BaseConfig's dict return type
+        optional_params: Mapping[str, object],
+        request_data: Mapping[str, object],
         api_base: str,
         api_key: str | None = None,
         model: str | None = None,
         stream: bool | None = None,
         fake_stream: bool | None = None,
-    ) -> tuple[dict, bytes | None]:  # mutable-ok: matches BaseConfig.sign_request's fixed override signature
+    ) -> tuple[dict[str, str], bytes | None]:  # mutable-ok: return type must match BaseConfig's dict-shaped contract
         if fake_stream is True:
             # /api/generate defaults to streaming when "stream" is absent from the body, but
             # the shared fake-stream handling drops the key instead of setting it False. Force

tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py

@@ -514,14 +514,6 @@ class TestOllamaTextCompletionResponseIterator:
 
 
 class TestOllamaFakeStreamActivation:
-    """Unit coverage for the #35711 fix's trigger condition: OllamaConfig.map_openai_params()
-    sets optional_params["fake_stream"] = True only when tool-call emulation is active
-    (functions_unsupported_model present, mirroring the pattern in
-    test_ollama_chat_transformation.py's assertions on get_optional_params()'s output)
-    and the caller actually requested streaming. Complements
-    TestOllamaFakeStreamToolCalls, which proves the full chain that this trigger feeds into.
-    """
-
     def _tools(self):
         return [
             {
@@ -568,21 +560,8 @@ class TestOllamaFakeStreamActivation:
 
 
 class TestOllamaFakeStreamToolCalls:
-    """Regression test for #35711 at the level the fix actually operates on.
-
-    OllamaTextCompletionResponseIterator.chunk_parser() never reconstructs tool calls
-    from streamed text and is not meant to: for ollama/ tool emulation (tools injected
-    into the prompt, never sent as a native `tools` field), OllamaConfig.map_openai_params()
-    sets fake_stream=True, which routes the request through a real non-streaming call to
-    Ollama, reuses the already-correct transform_response() reconstruction (see
-    test_transform_response_json_function_call above), and wraps the result as a single
-    fake stream chunk via MockResponseIterator. Only the HTTP boundary is mocked here;
-    the rest of litellm.completion()'s execution path runs for real, following the pattern
-    in test_vertex_gemma_transformation.py::test_acompletion_fake_streaming and
-    test_llm_http_handler.py::test_responses_handler_signs_after_fake_stream_prep_strips_stream.
-    """
-
     def test_tools_stream_true_reconstructs_tool_calls_via_fake_stream(self):
+        """Test that tools + stream=True routes through fake_stream and yields reconstructed tool_calls."""
         tool_call_json = {
             "name": "get_current_weather",
             "arguments": {"location": "San Francisco"},
@@ -595,68 +574,55 @@ class TestOllamaFakeStreamToolCalls:
             "prompt_eval_count": 42,
             "eval_count": 16,
         }
-        mock_ollama_response_bytes = json.dumps(mock_ollama_response).encode()
 
-        captured_requests = []
+        mock_client = MagicMock(spec=HTTPHandler)
+        mock_client.post.return_value = httpx.Response(
+            status_code=200,
+            content=json.dumps(mock_ollama_response).encode(),
+            request=httpx.Request("POST", "http://127.0.0.1:11434/api/generate"),
+        )
 
-        def _fake_post(
-            self, url, headers=None, data=None, timeout=None, stream=False, logging_obj=None, **kwargs
-        ):
-            request_body = json.loads(data) if isinstance(data, (str, bytes)) else {}
-            captured_requests.append({"url": url, "body": request_body})
-            return httpx.Response(
-                status_code=200,
-                content=mock_ollama_response_bytes,
-                request=httpx.Request("POST", url),
-            )
-
-        with patch.object(HTTPHandler, "post", _fake_post):
-            response = litellm.completion(
-                model="ollama/llama2",
-                api_base="http://127.0.0.1:11434",
-                messages=[
-                    {"role": "user", "content": "What is the weather in San Francisco?"}
-                ],
-                tools=[
-                    {
-                        "type": "function",
-                        "function": {
-                            "name": "get_current_weather",
-                            "description": "Get current weather.",
-                            "parameters": {
-                                "type": "object",
-                                "properties": {"location": {"type": "string"}},
-                                "required": ["location"],
-                            },
+        response = litellm.completion(
+            model="ollama/llama2",
+            api_base="http://127.0.0.1:11434",
+            messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "description": "Get current weather.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {"location": {"type": "string"}},
+                            "required": ["location"],
                         },
-                    }
-                ],
-                stream=True,
-                drop_params=True,
-            )
+                    },
+                }
+            ],
+            stream=True,
+            drop_params=True,
+            client=mock_client,
+        )
 
-            reassembled_content = ""
-            tool_calls_seen = []
-            finish_reasons = []
-            for chunk in response:
-                delta = chunk.choices[0].delta
-                if delta.content:
-                    reassembled_content += delta.content
-                if getattr(delta, "tool_calls", None):
-                    tool_calls_seen.extend(delta.tool_calls)
-                if chunk.choices[0].finish_reason:
-                    finish_reasons.append(chunk.choices[0].finish_reason)
+        reassembled_content = ""
+        tool_calls_seen = []
+        finish_reasons = []
+        for chunk in response:
+            delta = chunk.choices[0].delta
+            if delta.content:
+                reassembled_content += delta.content
+            if getattr(delta, "tool_calls", None):
+                tool_calls_seen.extend(delta.tool_calls)
+            if chunk.choices[0].finish_reason:
+                finish_reasons.append(chunk.choices[0].finish_reason)
 
-        assert len(captured_requests) == 1
-        request_body = captured_requests[0]["body"]
+        assert mock_client.post.call_count == 1
+        request_body = json.loads(mock_client.post.call_args.kwargs["data"])
 
-        # fake_stream engaged: the actual wire request to Ollama is non-streaming.
         assert request_body.get("stream") is False
-
-        # tools are emulated via the injected prompt, never sent as a native field.
         assert "tools" not in request_body
 
-        # the streamed response reconstructs the tool call, not raw JSON text.
         assert tool_calls_seen, "expected delta.tool_calls to be populated"
         assert tool_calls_seen[0]["function"]["name"] == "get_current_weather"
         assert json.loads(tool_calls_seen[0]["function"]["arguments"]) == {