diff --git a/garak/generators/ollama.py b/garak/generators/ollama.py index fb349dfef..6a74f2228 100644 --- a/garak/generators/ollama.py +++ b/garak/generators/ollama.py @@ -1,5 +1,6 @@ """Ollama interface""" +import logging from typing import List, Union import backoff @@ -23,11 +24,28 @@ class OllamaGenerator(Generator): """Interface for Ollama endpoints Model names can be passed in short form like "llama2" or specific versions or sizes like "gemma:7b" or "llama2:latest" + + suppressed_params (set[str], default empty): garak attribute names to omit from + the Ollama request `options` dict, regardless of whether the corresponding + attribute is set. Suppression is applied at request-assembly time, so it + overrides per-probe parameter mutation (such as promptinject's + _generator_precall_hook). Use garak attribute names (max_tokens, temperature, + top_k). + + Example garak.site.yaml config to suppress top_k:: + + plugins: + generators: + ollama: + OllamaGenerator: + suppressed_params: + - top_k """ DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "timeout": 30, # Add a timeout of 30 seconds. Ollama can tend to hang forever on failures, if this is not present "host": "127.0.0.1:11434", # The default host of an Ollama server. This can be overwritten with a passed config or generator config file. + "suppressed_params": set(), } active = True @@ -35,6 +53,17 @@ class OllamaGenerator(Generator): parallel_capable = False extra_dependency_names = ["ollama"] + # Maps garak attribute names to (Ollama options field name, coerce fn). + # Ollama takes generation parameters nested under `options`, and names the + # output cap `num_predict`, so the mapping is a translation rather than a + # passthrough. + _PARAM_MAP = { + "max_tokens": ("num_predict", int), + "temperature": ("temperature", float), + "top_k": ("top_k", int), + "seed": ("seed", int), + } + def __init__(self, name="", config_root=_config): super().__init__(name, config_root) # Sets the name and generations @@ -42,6 +71,30 @@ def __init__(self, name="", config_root=_config): self.host, timeout=self.timeout ) # Instantiates the client with the timeout + self.suppressed_params = set(self.suppressed_params) + for param in self.suppressed_params: + if param not in self._PARAM_MAP: + logging.warning( + f"suppressed_params entry '{param}' is not a known OllamaGenerator " + f"parameter. Valid keys are: {sorted(self._PARAM_MAP)}." + ) + + def _build_options(self): + """Assemble the Ollama `options` dict from configured generation params. + + Returns None when nothing is set, which the ollama client treats the same + as omitting the argument. + """ + options = {} + for attr, (api_field, coerce) in self._PARAM_MAP.items(): + if attr in self.suppressed_params: + continue + value = getattr(self, attr, None) + if value is None: + continue + options[api_field] = coerce(value) + return options or None + @backoff.on_exception( backoff.fibo, GeneratorBackoffTrigger, @@ -55,7 +108,11 @@ def _call_model( self, prompt: Conversation, generations_this_call: int = 1 ) -> List[Union[Message, None]]: try: - response = self.client.generate(self.name, prompt.last_message().text) + response = self.client.generate( + self.name, + prompt.last_message().text, + options=self._build_options(), + ) except Exception as e: if ( isinstance(e, self.ollama.ResponseError) and e.status_code == 404 @@ -94,6 +151,7 @@ def _call_model( response = self.client.chat( model=self.name, messages=messages, + options=self._build_options(), ) except Exception as e: if ( diff --git a/tests/generators/test_ollama.py b/tests/generators/test_ollama.py index 5d8d1dd36..4461a5885 100644 --- a/tests/generators/test_ollama.py +++ b/tests/generators/test_ollama.py @@ -1,4 +1,5 @@ import importlib +import json import pytest import respx import httpx @@ -214,3 +215,119 @@ def test_error_on_nonexistant_model_chat_mocked(respx_mock): with pytest.raises(ollama.ResponseError): conv = Conversation([Turn("user", Message("This shouldnt work"))]) gen.generate(conv) + + +@pytest.mark.skipif( + not all( + [ + importlib.util.find_spec(m) + for m in OllamaGeneratorChat.extra_dependency_names + ] + ), + reason="missing optional dependency", +) +@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"]) +def test_ollama_chat_forwards_generation_options(respx_mock): + mock_response = { + "model": "mistral", + "message": {"role": "assistant", "content": "Hello how are you?"}, + } + respx_mock.post("/api/chat").mock( + return_value=httpx.Response(200, json=mock_response) + ) + gen = OllamaGeneratorChat("mistral") + gen.max_tokens = 10 + gen.temperature = 0.1 + gen.top_k = 3 + gen.seed = 42 + conv = Conversation([Turn("user", Message("Bla bla"))]) + gen.generate(conv) + + sent = json.loads(respx_mock.calls.last.request.content) + assert sent["options"]["num_predict"] == 10 + assert sent["options"]["temperature"] == 0.1 + assert sent["options"]["top_k"] == 3 + assert sent["options"]["seed"] == 42 + + +@pytest.mark.skipif( + not all( + [importlib.util.find_spec(m) for m in OllamaGenerator.extra_dependency_names] + ), + reason="missing optional dependency", +) +@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"]) +def test_ollama_forwards_generation_options(respx_mock): + mock_response = {"model": "mistral", "response": "Hello how are you?"} + respx_mock.post("/api/generate").mock( + return_value=httpx.Response(200, json=mock_response) + ) + gen = OllamaGenerator("mistral") + gen.max_tokens = 10 + gen.temperature = 0.1 + gen.top_k = 3 + conv = Conversation([Turn("user", Message("Bla bla"))]) + gen.generate(conv) + + sent = json.loads(respx_mock.calls.last.request.content) + assert sent["options"]["num_predict"] == 10 + assert sent["options"]["temperature"] == 0.1 + assert sent["options"]["top_k"] == 3 + + +@pytest.mark.skipif( + not all( + [ + importlib.util.find_spec(m) + for m in OllamaGeneratorChat.extra_dependency_names + ] + ), + reason="missing optional dependency", +) +@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"]) +def test_ollama_chat_sends_default_max_tokens(respx_mock): + # max_tokens defaults to 150 in Generator.DEFAULT_PARAMS, so an unconfigured + # generator still caps output. Nothing else is set, so no other option is sent. + mock_response = { + "model": "mistral", + "message": {"role": "assistant", "content": "Hello how are you?"}, + } + respx_mock.post("/api/chat").mock( + return_value=httpx.Response(200, json=mock_response) + ) + gen = OllamaGeneratorChat("mistral") + conv = Conversation([Turn("user", Message("Bla bla"))]) + gen.generate(conv) + + sent = json.loads(respx_mock.calls.last.request.content) + assert sent["options"] == {"num_predict": gen.max_tokens} + + +@pytest.mark.skipif( + not all( + [ + importlib.util.find_spec(m) + for m in OllamaGeneratorChat.extra_dependency_names + ] + ), + reason="missing optional dependency", +) +@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"]) +def test_ollama_chat_honours_suppressed_params(respx_mock): + # suppression applies at request-assembly time, so it wins over a set attribute + mock_response = { + "model": "mistral", + "message": {"role": "assistant", "content": "Hello how are you?"}, + } + respx_mock.post("/api/chat").mock( + return_value=httpx.Response(200, json=mock_response) + ) + gen = OllamaGeneratorChat("mistral") + gen.temperature = 0.1 + gen.suppressed_params = {"temperature"} + conv = Conversation([Turn("user", Message("Bla bla"))]) + gen.generate(conv) + + sent = json.loads(respx_mock.calls.last.request.content) + assert "temperature" not in sent["options"] + assert sent["options"]["num_predict"] == gen.max_tokens