generators: make the ollama empty-response retry actually retry

garak/generators/ollama.py

@@ -117,7 +117,9 @@ def _build_options(self):
         giveup=_give_up,
     )
     @backoff.on_predicate(
-        backoff.fibo, lambda ans: ans == [None] or len(ans) == 0, max_tries=3
+        backoff.fibo,
+        lambda ans: not ans or ans[0] is None or not ans[0].text,
+        max_tries=3,
     )  # Ollama sometimes returns empty responses. Only 3 retries to not delay generations expecting empty responses too much
     def _call_model(
         self, prompt: Conversation, generations_this_call: int = 1
@@ -161,7 +163,9 @@ class OllamaGeneratorChat(OllamaGenerator):
         giveup=_give_up,
     )
     @backoff.on_predicate(
-        backoff.fibo, lambda ans: ans == [None] or len(ans) == 0, max_tries=3
+        backoff.fibo,
+        lambda ans: not ans or ans[0] is None or not ans[0].text,
+        max_tries=3,
     )  # Ollama sometimes returns empty responses. Only 3 retries to not delay generations expecting empty responses too much
     def _call_model(
         self, prompt: Conversation, generations_this_call: int = 1

tests/generators/test_ollama.py

@@ -192,6 +192,72 @@ def test_ollama_generation_chat_mocked(respx_mock):
     assert generation == [Message("Hello how are you?")]
 
 
+@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"])
+def test_empty_generation_is_retried(respx_mock):
+    """An empty generation is retried and a recovered response is returned."""
+    route = respx_mock.post("/api/generate").mock(
+        side_effect=[
+            httpx.Response(200, json={"model": "mistral", "response": None}),
+            httpx.Response(200, json={"model": "mistral", "response": "Recovered"}),
+        ]
+    )
+    gen = OllamaGenerator("mistral")
+    conv = Conversation([Turn("user", Message("Bla bla"))])
+    generation = gen.generate(conv)
+    assert generation == [
+        Message("Recovered")
+    ], "an empty generation must be retried and the recovered text surfaced"
+    assert (
+        route.call_count == 2
+    ), "the empty first response should have triggered exactly one retry"
+
+
+@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"])
+def test_empty_chat_generation_is_retried(respx_mock):
+    """An empty chat generation is retried and a recovered response is returned."""
+    route = respx_mock.post("/api/chat").mock(
+        side_effect=[
+            httpx.Response(
+                200,
+                json={"model": "mistral", "message": {"role": "assistant"}},
+            ),
+            httpx.Response(
+                200,
+                json={
+                    "model": "mistral",
+                    "message": {"role": "assistant", "content": "Recovered"},
+                },
+            ),
+        ]
+    )
+    gen = OllamaGeneratorChat("mistral")
+    conv = Conversation([Turn("user", Message("Bla bla"))])
+    generation = gen.generate(conv)
+    assert generation == [
+        Message("Recovered")
+    ], "an empty chat generation must be retried and the recovered text surfaced"
+    assert (
+        route.call_count == 2
+    ), "the empty first chat response should have triggered exactly one retry"
+
+
+@pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"])
+def test_content_generation_is_not_retried(respx_mock):
+    """A response with content goes through on the first attempt."""
+    route = respx_mock.post("/api/generate").mock(
+        return_value=httpx.Response(200, json={"model": "mistral", "response": "Hi"})
+    )
+    gen = OllamaGenerator("mistral")
+    conv = Conversation([Turn("user", Message("Bla bla"))])
+    generation = gen.generate(conv)
+    assert generation == [
+        Message("Hi")
+    ], "the content-bearing response is returned as is"
+    assert (
+        route.call_count == 1
+    ), "a generation carrying content must not trigger retries"
+
+
 @pytest.mark.respx(base_url="http://" + OllamaGenerator.DEFAULT_PARAMS["host"])
 def test_error_on_nonexistant_model_mocked(respx_mock):
     mock_response = {"error": "No such model"}