"""Tests for llm capability detection and feature fallbacks.""" from __future__ import annotations from types import SimpleNamespace from typing import Any import pytest from ollama import Message from ollama._types import Image as OllamaImage from ollama_web.llm import ( _is_image_message, _remove_images_with_notice, _token_counts, stream_chat_with_tools, ) from ollama_web.tools.registry import ToolArtifact, ToolExecutionResult, ToolRegistry def _image_message() -> Message: return Message( role="user", content="hello", images=[OllamaImage(value=b"fake-image-bytes")], ) def test_is_image_message_detects_images() -> None: assert _is_image_message(_image_message()) is True def test_is_image_message_without_images() -> None: assert _is_image_message(Message(role="user", content="hello")) is False def test_remove_images_adds_localized_notice_without_base64() -> None: converted = _remove_images_with_notice(_image_message(), "ja") assert "画像入力に対応していません" in str(converted.content) assert "ZmFrZS1pbWFnZS1ieXRlcw" not in str(converted.content) assert not getattr(converted, "images", None) 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() def lookup(query: str) -> str: return query registry.register(lookup) 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) events = list( stream_chat_with_tools( [Message(role="user", content="hi")], model="test-model", think=True, capabilities=set(), registry=registry, ) ) assert client.calls[0]["tools"] == [] assert client.calls[0]["think"] is None assert client.calls[0]["stream"] is True assert events[-1] == {"type": "done"} @pytest.mark.parametrize( ("capabilities", "expect_images", "notice"), [ ({"tools", "vision"}, True, None), ({"tools"}, False, "画像の内容はモデルへ提供されていません"), ], ) def test_mcp_image_is_only_sent_to_vision_models( monkeypatch: pytest.MonkeyPatch, capabilities: set[str], expect_images: bool, notice: str | None, ) -> None: from ollama_web import llm registry = ToolRegistry() def executor(full_name: str, arguments: dict[str, object]) -> ToolExecutionResult: return ToolExecutionResult( text="tool text", artifacts=[ ToolArtifact( kind="image", mime="image/png", name="generated.png", data=b"fake-image-bytes", file_id="a" * 32, ) ], ) registry.register_mcp_tool( "mcp__test__image", {"type": "function", "function": {"name": "mcp__test__image"}}, 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": "mcp__test__image", "arguments": {}, } } ], } ) ] return [SimpleNamespace(message={"content": "done"})] client = FakeClient() monkeypatch.setattr(llm, "get_client", lambda _host=None: client) monkeypatch.setattr(llm, "resize_image", lambda data, name="": data) events = list( llm.stream_chat_with_tools( [Message(role="user", content="generate")], model="test-model", registry=registry, capabilities=capabilities, language="ja", ) ) tool_message = next( message for message in reversed(client.calls[1]["messages"]) if getattr(message, "role", None) == "tool" ) assert bool(getattr(tool_message, "images", None)) is expect_images if notice is not None: assert notice in str(tool_message.content) assert any( event.get("media", [{}])[0].get("id") == "a" * 32 for event in events if event.get("media") ) def test_mcp_audio_is_never_sent_to_ollama( monkeypatch: pytest.MonkeyPatch, ) -> None: from ollama_web import llm registry = ToolRegistry() def executor(full_name: str, arguments: dict[str, object]) -> ToolExecutionResult: return ToolExecutionResult( text="music generated", artifacts=[ ToolArtifact( kind="audio", mime="audio/mpeg", name="generated.mp3", data=b"audio-bytes", file_id="b" * 32, ) ], ) registry.register_mcp_tool( "mcp__test__audio", {"type": "function", "function": {"name": "mcp__test__audio"}}, 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": "mcp__test__audio", "arguments": {}, } } ], } ) ] return [SimpleNamespace(message={"content": "done"})] client = FakeClient() monkeypatch.setattr(llm, "get_client", lambda _host=None: client) list( llm.stream_chat_with_tools( [Message(role="user", content="generate")], model="test-model", registry=registry, capabilities={"tools", "vision"}, language="ja", ) ) tool_message = next( message for message in reversed(client.calls[1]["messages"]) if getattr(message, "role", None) == "tool" ) assert not getattr(tool_message, "images", None) assert "音声データそのものはモデルへ提供されていません" in str(tool_message.content) assert "YXVkaW8tYnl0ZXM" not in str(tool_message.content) def test_token_counts_extracts_from_dict() -> None: assert _token_counts({"prompt_eval_count": 100, "eval_count": 50}) == { "prompt_eval_count": 100, "eval_count": 50, } def test_token_counts_extracts_from_object() -> None: class FakeResp: prompt_eval_count = 200 eval_count = 30 assert _token_counts(FakeResp()) == { "prompt_eval_count": 200, "eval_count": 30, } def test_token_counts_returns_none_when_missing() -> None: assert _token_counts({}) is None assert _token_counts(None) is None def test_stream_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) -> list[SimpleNamespace]: self.calls.append(kwargs) if len(self.calls) == 1: 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.stream_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 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"} 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"} def test_paged_source_limit_emits_incomplete_warning( monkeypatch: pytest.MonkeyPatch, ) -> None: from ollama_web import llm registry = ToolRegistry() def scrape_url(url: str) -> str: """Return a paged result.""" return "body\n---page-info---\nhas_more: true\nnext_token: token" registry.register(scrape_url) 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": "scrape_url", "arguments": {"url": "https://example.com"}, } } ], } ) ] return [SimpleNamespace(message={"content": "partial 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="read all")], model="test-model", registry=registry, capabilities={"tools"}, language="en", ) ) assert any(event.get("type") == "warning" for event in events) final_instruction = str(client.calls[-1]["messages"][-1].content) assert "unread content" in final_instruction def test_paged_tool_is_not_executed_when_context_budget_is_exhausted( monkeypatch: pytest.MonkeyPatch, ) -> None: from ollama_web import llm executed: list[str] = [] registry = ToolRegistry() def pdf_to_text(url_or_path: str) -> str: """Read a PDF.""" executed.append(url_or_path) return "should not execute" registry.register(pdf_to_text) class FakeClient: def __init__(self) -> None: self.calls = 0 def chat(self, **kwargs: Any) -> list[SimpleNamespace]: self.calls += 1 if self.calls == 1: return [ SimpleNamespace( message={ "content": "", "tool_calls": [ { "function": { "name": "pdf_to_text", "arguments": {"url_or_path": "https://example.com/a.pdf"}, } } ], } ) ] return [SimpleNamespace(message={"content": "partial"})] monkeypatch.setattr(llm, "get_client", lambda _host=None: FakeClient()) monkeypatch.setattr(llm, "MAX_TOOL_CONTEXT_CHARS", 1000) monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1) events = list( llm.stream_chat_with_tools( [Message(role="user", content="read all")], model="test-model", registry=registry, capabilities={"tools"}, language="en", ) ) assert executed == [] assert any("budget exhausted" in event.get("result", "") for event in events)