1 files changed, 53 insertions, 48 deletions

diff --git a/tests/test_llm.py b/tests/test_llm.py
index edd9d61..0e1b146 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -9,13 +9,11 @@ from typing import Any
 import pytest
 from ollama import Message
 from ollama._types import Image as OllamaImage
-from ollama._types import ResponseError
 
 from ollama_web.llm import (
     _embed_images_in_content,
     _is_image_message,
     _token_counts,
-    chat_with_tools,
     stream_chat_with_tools,
 )
 from ollama_web.tools.registry import ToolRegistry
@@ -44,46 +42,44 @@ def test_embed_images_in_content_adds_base64_markdown() -> None:
     assert not getattr(embedded, "images", None)
 
 
-def test_chat_with_tools_skips_tools_for_non_tool_model() -> None:
-    """Tools are not passed when the model does not advertise tool support."""
-    messages = [Message(role="user", content="hi")]
-    # Use an empty registry so any tool call would fail if requested.
+def test_stream_chat_with_tools_disables_unsupported_capabilities(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Unsupported tools and thinking options are not sent to Ollama."""
+    from ollama_web import llm
+
     registry = ToolRegistry()
 
-    # Without tool capability, the iterator should finish immediately with no
-    # actual network request. We only verify it accepts the call and does not
-    # raise due to tool-related setup.
-    gen = chat_with_tools(
-        messages,
-        model="nonexistent-model",
-        capabilities=set(),
-        registry=registry,
-    )
-    events = list(gen)
-    # We expect either an error (model not found) or a done event.
-    assert any(e.get("type") in {"error", "done"} for e in events)
+    def lookup(query: str) -> str:
+        return query
 
+    registry.register(lookup)
 
-def test_chat_with_tools_forces_think_none_for_non_thinking_model() -> None:
-    messages = [Message(role="user", content="hi")]
-    gen = chat_with_tools(
-        messages,
-        model="nonexistent-model",
-        capabilities=set(),
-    )
-    events = list(gen)
-    assert any(e.get("type") in {"error", "done"} for e in events)
+    class FakeClient:
+        def __init__(self) -> None:
+            self.calls: list[dict[str, Any]] = []
 
+        def chat(self, **kwargs: Any) -> list[SimpleNamespace]:
+            self.calls.append(kwargs)
+            return [SimpleNamespace(message={"content": "done"})]
+
+    client = FakeClient()
+    monkeypatch.setattr(llm, "get_client", lambda _host=None: client)
 
-def test_stream_chat_with_tools_accepts_capability_set() -> None:
-    messages = [Message(role="user", content="hi")]
-    gen = stream_chat_with_tools(
-        messages,
-        model="nonexistent-model",
-        capabilities={"tools", "thinking"},
+    events = list(
+        stream_chat_with_tools(
+            [Message(role="user", content="hi")],
+            model="test-model",
+            think=True,
+            capabilities=set(),
+            registry=registry,
+        )
     )
-    with pytest.raises(ResponseError):
-        list(gen)
+
+    assert client.calls[0]["tools"] == []
+    assert client.calls[0]["think"] is None
+    assert client.calls[0]["stream"] is True
+    assert events[-1] == {"type": "done"}
 
 
 def test_token_counts_extracts_from_dict() -> None:
@@ -109,7 +105,7 @@ def test_token_counts_returns_none_when_missing() -> None:
     assert _token_counts(None) is None
 
 
-def test_chat_with_tools_disables_tools_for_final_answer_after_limit(
+def test_stream_chat_with_tools_disables_tools_for_final_answer_after_limit(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     from ollama_web import llm
@@ -128,26 +124,33 @@ def test_chat_with_tools_disables_tools_for_final_answer_after_limit(
         def __init__(self) -> None:
             self.calls: list[dict[str, Any]] = []
 
-        def chat(self, **kwargs: Any) -> dict[str, Any]:
+        def chat(self, **kwargs: Any) -> list[SimpleNamespace]:
             self.calls.append(kwargs)
             if len(self.calls) == 1:
-                return {
-                    "message": {
-                        "role": "assistant",
-                        "content": "",
-                        "tool_calls": [
-                            {"function": {"name": "lookup", "arguments": {"query": "test"}}}
-                        ],
-                    }
-                }
-            return {"message": {"role": "assistant", "content": "final answer"}}
+                return [
+                    SimpleNamespace(
+                        message={
+                            "role": "assistant",
+                            "content": "",
+                            "tool_calls": [
+                                {
+                                    "function": {
+                                        "name": "lookup",
+                                        "arguments": {"query": "test"},
+                                    }
+                                }
+                            ],
+                        }
+                    )
+                ]
+            return [SimpleNamespace(message={"content": "final answer"})]
 
     client = FakeClient()
     monkeypatch.setattr(llm, "get_client", lambda _host=None: client)
     monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1)
 
     events = list(
-        llm.chat_with_tools(
+        llm.stream_chat_with_tools(
             [Message(role="user", content="research this")],
             model="test-model",
             registry=registry,
@@ -160,6 +163,8 @@ def test_chat_with_tools_disables_tools_for_final_answer_after_limit(
     assert len(client.calls) == 2
     assert client.calls[0]["tools"]
     assert client.calls[1]["tools"] == []
+    assert client.calls[0]["stream"] is True
+    assert client.calls[1]["stream"] is True
     assert any(event.get("type") == "status" for event in events)
     assert any(event == {"type": "delta", "content": "final answer"} for event in events)
     assert events[-1] == {"type": "done"}