tests/test_ollama.py

import json
from pathlib import Path

from PIL import Image

from adam.ollama import OllamaClient


def test_qwen3_chat_has_a_larger_reply_budget_for_reasoning() -> None:
    client = OllamaClient("http://localhost:11434", "qwen3:4b")

    assert client._chat_system("You are ADAM.") == "You are ADAM."
    assert client._num_predict(180) == 1024


def test_other_models_keep_their_original_system_prompt() -> None:
    client = OllamaClient("http://localhost:11434", "qwen2.5:1.5b")

    assert client._chat_system("You are ADAM.") == "You are ADAM."
    assert client._num_predict(180) == 180


def test_automatic_chat_budget_stays_under_the_configured_limit() -> None:
    client = OllamaClient("http://localhost:11434", "qwen3:4b", chat_max_tokens=1536)

    assert client._num_predict(180, prompt="hello") == 256
    assert client._num_predict(180, prompt="Explain how this training plan works") == 768
    assert client._num_predict(180, prompt="Describe this image", image_count=1) == 384
    assert client._num_predict(300, chat=False) == 1024


def test_explicit_detailed_mode_uses_the_configured_limit() -> None:
    client = OllamaClient(
        "http://localhost:11434", "qwen3:4b", chat_max_tokens=1536,
        chat_response_length="detailed",
    )

    assert client._num_predict(180, prompt="hello") == 1536


def test_streaming_vision_request_attaches_a_resized_image(tmp_path: Path, monkeypatch) -> None:
    image = tmp_path / "reference.png"
    Image.new("RGB", (2400, 1200), (24, 48, 72)).save(image)
    captured = {}

    class Response:
        status = 200

        def __enter__(self): return self
        def __exit__(self, *_args): return None
        def __iter__(self):
            return iter([b'{"response":"A clear blue image.","done":true}\n'])

    def open_request(request, timeout):
        captured["payload"] = json.loads(request.data.decode("utf-8"))
        captured["timeout"] = timeout
        return Response()

    monkeypatch.setattr("adam.ollama.urllib.request.urlopen", open_request)
    chunks = []
    client = OllamaClient("http://localhost:11434", "qwen3-vl:4b", chat_max_tokens=1024)

    assert client.generate_text_stream("You are ADAM.", "Describe this image", chunks.append, [image]) == "A clear blue image."
    assert chunks == ["A clear blue image."]
    assert captured["payload"]["options"]["num_predict"] == 384
    assert len(captured["payload"]["images"]) == 1