Add Retries for LLM-Guided Metrics (#728)

Co-authored-by: b.nativi <[email protected]>

Author: bnativi

api/tests/functional-tests/backend/core/test_llm_clients.py

@@ -1,6 +1,6 @@
 import datetime
 import os
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 from mistralai.models import (
@@ -1320,6 +1320,109 @@ def _return_invalid4_toxicity_response(*args, **kwargs):
         client.toxicity("some text")
 
 
+def test_LLMClient_retries(monkeypatch):
+    """
+    Test the retry functionality for structuring LLM API calls.
+    """
+
+    def _return_valid_summary_coherence_response(*args, **kwargs):
+        return "5"
+
+    errors = ["The score is 5."] * 3 + ["5"]
+
+    def _return_invalid_summary_coherence_response(*args, **kwargs):
+        return "The score is 5."
+
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.LLMClient.__call__",
+        _return_valid_summary_coherence_response,
+    )
+
+    # Test with retries=None
+    client = LLMClient(api_key=None, model_name="model_name", retries=None)
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    # Test with retries=0
+    client = LLMClient(api_key=None, model_name="model_name", retries=0)
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    # Test with retries=3 and valid response
+    client = LLMClient(api_key=None, model_name="model_name", retries=3)
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    # mock_method returns a bad response three times but on the fourth call returns a valid response.
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.LLMClient.__call__",
+        Mock(side_effect=errors),
+    )
+    client = LLMClient(api_key=None, model_name="model_name", retries=3)
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    # Test with retries=2 and invalid response
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.LLMClient.__call__",
+        Mock(side_effect=errors),
+    )
+    with pytest.raises(InvalidLLMResponseError):
+        client = LLMClient(api_key=None, model_name="model_name", retries=2)
+        client.summary_coherence("some text", "some summary")
+
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.LLMClient.__call__",
+        _return_invalid_summary_coherence_response,
+    )
+
+    # Test with retries=None and invalid response
+    with pytest.raises(InvalidLLMResponseError):
+        client = LLMClient(api_key=None, model_name="model_name", retries=None)
+        client.summary_coherence("some text", "some summary")
+
+    # Test with retries=3 and invalid response
+    with pytest.raises(InvalidLLMResponseError):
+        client = LLMClient(api_key=None, model_name="model_name", retries=3)
+        client.summary_coherence("some text", "some summary")
+
+    # Test WrappedOpenAIClient
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.WrappedOpenAIClient.__call__",
+        Mock(side_effect=errors),
+    )
+    client = WrappedOpenAIClient(
+        api_key=None, model_name="model_name", retries=3
+    )
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    with pytest.raises(InvalidLLMResponseError):
+        monkeypatch.setattr(
+            "valor_api.backend.core.llm_clients.WrappedOpenAIClient.__call__",
+            Mock(side_effect=errors),
+        )
+        client = WrappedOpenAIClient(
+            api_key=None, model_name="model_name", retries=2
+        )
+        client.summary_coherence("some text", "some summary")
+
+    # Test WrappedMistralAIClient
+    monkeypatch.setattr(
+        "valor_api.backend.core.llm_clients.WrappedMistralAIClient.__call__",
+        Mock(side_effect=errors),
+    )
+    client = WrappedMistralAIClient(
+        api_key=None, model_name="model_name", retries=3
+    )
+    assert 5 == client.summary_coherence("some text", "some summary")
+
+    with pytest.raises(InvalidLLMResponseError):
+        monkeypatch.setattr(
+            "valor_api.backend.core.llm_clients.WrappedMistralAIClient.__call__",
+            Mock(side_effect=errors),
+        )
+        client = WrappedMistralAIClient(
+            api_key=None, model_name="model_name", retries=2
+        )
+        client.summary_coherence("some text", "some summary")
+
+
 def test_WrappedOpenAIClient():
     def _create_bad_request(model, messages, seed) -> ChatCompletion:
         raise ValueError