From be11db60066c849c59f7774ce633e7100b4654be Mon Sep 17 00:00:00 2001
From: Carles Onielfa <carlesop@nuclia.com>
Date: Thu, 19 Dec 2024 08:47:00 +0100
Subject: [PATCH] Add NUA REMi E2E

---
 nua/e2e/regional/test_predict.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/nua/e2e/regional/test_predict.py b/nua/e2e/regional/test_predict.py
index f5d3d1c..ac9db58 100644
--- a/nua/e2e/regional/test_predict.py
+++ b/nua/e2e/regional/test_predict.py
@@ -2,6 +2,7 @@
 from nuclia.sdk.predict import AsyncNucliaPredict
 
 from regional.models import ALL_ENCODERS, ALL_LLMS
+from nuclia_models.predict.remi import RemiRequest
 
 
 @pytest.mark.asyncio_cooperative
@@ -55,3 +56,27 @@ async def test_predict_rephrase(nua_config, model):
     # TODO: Test that custom rephrase prompt works once SDK supports it
     rephrased = await np.rephrase(question="Barcelona best coffe", model=model)
     assert rephrased != "Barcelona best coffe" and rephrased != ""
+
+
+@pytest.mark.asyncio_cooperative
+async def test_predict_remi(nua_config):
+    # Check that rephrase is working for all models
+    np = AsyncNucliaPredict()
+    results = await np.remi(
+        RemiRequest(
+            user_id="NUA E2E",
+            question="What is the capital of France?",
+            answer="Paris is the capital of france!",
+            contexts=[
+                "Paris is the capital of France.",
+                "Berlin is the capital of Germany.",
+            ],
+        )
+    )
+    assert results.answer_relevance.score >= 4
+
+    assert results.context_relevance[0] >= 4
+    assert results.groundedness[0] >= 4
+
+    assert results.context_relevance[1] < 2
+    assert results.groundedness[1] < 2