diff options
| -rw-r--r-- | src/ollama_web/llm.py | 6 | ||||
| -rw-r--r-- | tests/test_llm.py | 126 |
2 files changed, 129 insertions, 3 deletions
diff --git a/src/ollama_web/llm.py b/src/ollama_web/llm.py index 2839051..4951b36 100644 --- a/src/ollama_web/llm.py +++ b/src/ollama_web/llm.py @@ -326,7 +326,7 @@ def chat_with_tools( resp = cast(Any, raw_client).chat( model=model, messages=messages, - tools=effective_tools, + tools=[], think=effective_think, options=options or None, stream=False, @@ -596,7 +596,7 @@ def stream_chat_with_tools( stream = cast(Any, raw_client).chat( model=model, messages=messages, - tools=effective_tools, + tools=[], think=effective_think, options=options or None, stream=True, @@ -716,4 +716,4 @@ def _assistant_message_chunk(full_message: dict[str, Any]) -> Message | None: Message.ToolCall(function=Message.ToolCall.Function(name=name, arguments=args)) ) kwargs["tool_calls"] = tc_list - return Message(**kwargs)
\ No newline at end of file + return Message(**kwargs) diff --git a/tests/test_llm.py b/tests/test_llm.py index b4e68ad..edd9d61 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -3,6 +3,8 @@ from __future__ import annotations import base64 +from types import SimpleNamespace +from typing import Any import pytest from ollama import Message @@ -105,3 +107,127 @@ def test_token_counts_extracts_from_object() -> None: def test_token_counts_returns_none_when_missing() -> None: assert _token_counts({}) is None assert _token_counts(None) is None + + +def test_chat_with_tools_disables_tools_for_final_answer_after_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from ollama_web import llm + + executed: list[str] = [] + + def lookup(query: str) -> str: + """Return a test lookup result.""" + executed.append(query) + return "lookup result" + + registry = ToolRegistry() + registry.register(lookup) + + class FakeClient: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def chat(self, **kwargs: Any) -> dict[str, Any]: + 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"}} + + client = FakeClient() + monkeypatch.setattr(llm, "get_client", lambda _host=None: client) + monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1) + + events = list( + llm.chat_with_tools( + [Message(role="user", content="research this")], + model="test-model", + registry=registry, + capabilities={"tools"}, + language="en", + ) + ) + + assert executed == ["test"] + assert len(client.calls) == 2 + assert client.calls[0]["tools"] + assert client.calls[1]["tools"] == [] + 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"} + + +def test_stream_chat_with_tools_disables_mcp_tools_for_final_answer_after_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from ollama_web import llm + + tool_name = "mcp__test__lookup" + executed: list[tuple[str, dict[str, object] | str]] = [] + registry = ToolRegistry() + + def executor(full_name: str, arguments: dict[str, object] | str = "") -> str: + executed.append((full_name, arguments)) + return "mcp lookup result" + + registry.register_mcp_tool( + tool_name, + {"type": "function", "function": {"name": tool_name}}, + executor, + ) + + class FakeClient: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def chat(self, **kwargs: Any) -> list[SimpleNamespace]: + self.calls.append(kwargs) + if len(self.calls) == 1: + return [ + SimpleNamespace( + message={ + "content": "", + "tool_calls": [ + { + "function": { + "name": tool_name, + "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.stream_chat_with_tools( + [Message(role="user", content="research this")], + model="test-model", + registry=registry, + capabilities={"tools"}, + language="en", + ) + ) + + assert executed == [(tool_name, {"query": "test"})] + 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"} |
