Forward generation parameters to Ollama via options dict

garak/generators/ollama.py

@@ -1,5 +1,6 @@
 """Ollama interface"""
 
+import logging
 from typing import List, Union
 
 import backoff
@@ -23,25 +24,77 @@ 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
     generator_family_name = "Ollama"
     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
 
         self.client = self.ollama.Client(
             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 (

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