fix(reranker): support nested llm config in LLMReranker for non-OpenAI providers (#4405)

tests/rerankers/test_llm_reranker_rerank.py

@@ -0,0 +1,125 @@
+import pytest
+
+from mem0.reranker.llm_reranker import LLMReranker
+
+
+class TestExtractScore:
+    @pytest.fixture
+    def reranker(self, mock_llm):
+        return LLMReranker({"provider": "openai"})
+
+    @pytest.mark.parametrize(
+        "text,expected",
+        [
+            ("0.85", 0.85),
+            ("0.0", 0.0),
+            ("1.0", 1.0),
+            ("The score is 0.72.", 0.72),
+            ("Score: 0.9 out of 1.0", 0.9),
+        ],
+    )
+    def test_valid_scores(self, reranker, text, expected):
+        assert reranker._extract_score(text) == expected
+
+    def test_no_score_returns_fallback(self, reranker):
+        assert reranker._extract_score("no numbers here") == 0.5
+
+    def test_clamps_to_1(self, reranker):
+        assert reranker._extract_score("1.0") == 1.0
+
+
+class TestRerank:
+    def test_empty_documents(self, mock_llm):
+        reranker = LLMReranker({"provider": "openai"})
+        result = reranker.rerank("query", [])
+        assert result == []
+
+    def test_documents_sorted_by_score_descending(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.side_effect = ["0.3", "0.9", "0.6"]
+
+        reranker = LLMReranker({"provider": "openai"})
+        docs = [
+            {"memory": "low relevance"},
+            {"memory": "high relevance"},
+            {"memory": "mid relevance"},
+        ]
+
+        result = reranker.rerank("test query", docs)
+
+        assert len(result) == 3
+        assert result[0]["rerank_score"] == 0.9
+        assert result[1]["rerank_score"] == 0.6
+        assert result[2]["rerank_score"] == 0.3
+
+    def test_top_k_limits_results(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.side_effect = ["0.9", "0.5", "0.1"]
+
+        reranker = LLMReranker({"provider": "openai"})
+        docs = [{"memory": f"doc{i}"} for i in range(3)]
+
+        result = reranker.rerank("query", docs, top_k=2)
+        assert len(result) == 2
+
+    def test_config_top_k_used_when_arg_not_provided(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.side_effect = ["0.9", "0.5", "0.1"]
+
+        reranker = LLMReranker({"provider": "openai", "top_k": 1})
+        docs = [{"memory": f"doc{i}"} for i in range(3)]
+
+        result = reranker.rerank("query", docs)
+        assert len(result) == 1
+
+    def test_text_field_extraction(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.return_value = "0.8"
+
+        reranker = LLMReranker({"provider": "openai"})
+        reranker.rerank("query", [{"text": "some text"}])
+
+        prompt_sent = mock_llm_instance.generate_response.call_args[1]["messages"][0]["content"]
+        assert "some text" in prompt_sent
+
+    def test_content_field_extraction(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.return_value = "0.8"
+
+        reranker = LLMReranker({"provider": "openai"})
+        reranker.rerank("query", [{"content": "some content"}])
+
+        prompt_sent = mock_llm_instance.generate_response.call_args[1]["messages"][0]["content"]
+        assert "some content" in prompt_sent
+
+    def test_fallback_score_on_llm_error(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.side_effect = RuntimeError("API error")
+
+        reranker = LLMReranker({"provider": "openai"})
+        result = reranker.rerank("query", [{"memory": "doc"}])
+
+        assert len(result) == 1
+        assert result[0]["rerank_score"] == 0.5
+
+    def test_custom_scoring_prompt(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.return_value = "0.7"
+
+        custom_prompt = "Rate this: query={query} doc={document}"
+        reranker = LLMReranker({"provider": "openai", "scoring_prompt": custom_prompt})
+        reranker.rerank("my query", [{"memory": "my doc"}])
+
+        prompt_sent = mock_llm_instance.generate_response.call_args[1]["messages"][0]["content"]
+        assert prompt_sent == "Rate this: query=my query doc=my doc"
+
+    def test_original_doc_not_mutated(self, mock_llm):
+        _, mock_llm_instance = mock_llm
+        mock_llm_instance.generate_response.return_value = "0.8"
+
+        reranker = LLMReranker({"provider": "openai"})
+        original_doc = {"memory": "test", "id": "123"}
+        result = reranker.rerank("query", [original_doc])
+
+        assert "rerank_score" not in original_doc
+        assert "rerank_score" in result[0]