chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,481 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from jinja2 import TemplateSyntaxError
|
||||
|
||||
from haystack import Document
|
||||
from haystack.components.generators.chat.openai import OpenAIChatGenerator
|
||||
from haystack.components.rankers.llm_ranker import DEFAULT_PROMPT_TEMPLATE, LLMRanker
|
||||
from haystack.dataclasses import ChatMessage
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_generator():
|
||||
return Mock(spec=OpenAIChatGenerator)
|
||||
|
||||
|
||||
def test_init_invalid_top_k():
|
||||
with pytest.raises(ValueError, match="top_k must be > 0"):
|
||||
LLMRanker(top_k=0)
|
||||
|
||||
|
||||
def test_init_default_generator(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
||||
ranker = LLMRanker()
|
||||
|
||||
assert ranker.top_k == 10
|
||||
assert ranker.raise_on_failure is False
|
||||
assert ranker.prompt == DEFAULT_PROMPT_TEMPLATE
|
||||
assert isinstance(ranker._chat_generator, OpenAIChatGenerator)
|
||||
assert ranker._chat_generator.model == "gpt-4.1-mini"
|
||||
assert ranker._prompt_builder is not None
|
||||
|
||||
|
||||
def test_init_custom_generator(mock_chat_generator):
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=5, raise_on_failure=True)
|
||||
|
||||
assert ranker._chat_generator is mock_chat_generator
|
||||
assert ranker.top_k == 5
|
||||
assert ranker.raise_on_failure is True
|
||||
|
||||
|
||||
def test_to_dict(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
||||
chat_generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0.5})
|
||||
ranker = LLMRanker(
|
||||
chat_generator=chat_generator,
|
||||
prompt="Rank {{ documents|length }} docs for {{ query }}",
|
||||
top_k=3,
|
||||
raise_on_failure=True,
|
||||
)
|
||||
|
||||
assert ranker.to_dict() == {
|
||||
"type": "haystack.components.rankers.llm_ranker.LLMRanker",
|
||||
"init_parameters": {
|
||||
"chat_generator": chat_generator.to_dict(),
|
||||
"prompt": "Rank {{ documents|length }} docs for {{ query }}",
|
||||
"top_k": 3,
|
||||
"raise_on_failure": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_from_dict(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
||||
chat_generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0.5})
|
||||
data = {
|
||||
"type": "haystack.components.rankers.llm_ranker.LLMRanker",
|
||||
"init_parameters": {
|
||||
"chat_generator": chat_generator.to_dict(),
|
||||
"prompt": "Rank {{ documents|length }} docs for {{ query }}",
|
||||
"top_k": 3,
|
||||
"raise_on_failure": True,
|
||||
},
|
||||
}
|
||||
|
||||
ranker = LLMRanker.from_dict(data)
|
||||
|
||||
assert ranker.top_k == 3
|
||||
assert ranker.raise_on_failure is True
|
||||
assert ranker.prompt == "Rank {{ documents|length }} docs for {{ query }}"
|
||||
assert ranker._chat_generator.to_dict() == chat_generator.to_dict()
|
||||
|
||||
|
||||
def test_run_invalid_runtime_top_k(mock_chat_generator):
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
with pytest.raises(ValueError, match="top_k must be > 0"):
|
||||
ranker.run(query="test", documents=[Document(content="doc")], top_k=0)
|
||||
|
||||
|
||||
def test_run_empty_documents(mock_chat_generator):
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
assert ranker.run(query="test", documents=[]) == {"documents": []}
|
||||
|
||||
|
||||
def test_run_whitespace_query_returns_fallback(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
|
||||
|
||||
result = ranker.run(query=" ", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
mock_chat_generator.run.assert_not_called()
|
||||
|
||||
|
||||
def test_run_successful_ranking(mock_chat_generator):
|
||||
documents = [
|
||||
Document(id="1", content="first"),
|
||||
Document(id="2", content="second"),
|
||||
Document(id="3", content="third"),
|
||||
]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2", "1"]
|
||||
|
||||
|
||||
def test_run_returns_only_documents_listed_by_the_llm(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}]}')]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2"]
|
||||
|
||||
|
||||
def test_run_runtime_top_k_overrides_instance_top_k(mock_chat_generator):
|
||||
documents = [
|
||||
Document(id="doc_1", content="first"),
|
||||
Document(id="doc_2", content="second"),
|
||||
Document(id="doc_3", content="third"),
|
||||
]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 3}, {"index": 2}, {"index": 1}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=3)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents, top_k=1)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["doc_3"]
|
||||
|
||||
|
||||
def test_run_ignores_out_of_range_indices(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 99}, {"index": 2}, {"index": 1}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2", "1"]
|
||||
|
||||
|
||||
def test_run_empty_ranking_result_returns_empty_documents(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": []}')]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": []}
|
||||
|
||||
|
||||
def test_run_invalid_json_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1, raise_on_failure=False)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_run_invalid_json_raises(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, raise_on_failure=True)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ranker.run(query="test query", documents=documents)
|
||||
|
||||
|
||||
def test_run_generator_exception_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.side_effect = RuntimeError("generator failed")
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_run_generator_exception_raises(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first")]
|
||||
mock_chat_generator.run.side_effect = RuntimeError("generator failed")
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, raise_on_failure=True)
|
||||
|
||||
with pytest.raises(RuntimeError, match="generator failed"):
|
||||
ranker.run(query="test query", documents=documents)
|
||||
|
||||
|
||||
def test_run_no_replies_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": []}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_run_reply_without_text_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant(tool_calls=[])]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_run_no_valid_document_indices_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 0}, {"index": 3}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_run_deduplicates_documents_before_ranking(mock_chat_generator):
|
||||
documents = [
|
||||
Document(id="duplicate", content="keep me", score=0.9),
|
||||
Document(id="duplicate", content="drop me", score=0.1),
|
||||
Document(id="unique", content="unique", score=0.2),
|
||||
]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert [document.content for document in result["documents"]] == ["unique", "keep me"]
|
||||
|
||||
|
||||
def test_run_preserves_duplicate_indices(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 2}, {"index": 1}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2", "2", "1"]
|
||||
|
||||
|
||||
def test_run_numeric_string_index_is_accepted(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": "2"}]}')]}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": [documents[1]]}
|
||||
|
||||
|
||||
def test_run_invalid_index_type_falls_back(mock_chat_generator):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": "invalid"}]}')]
|
||||
}
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
|
||||
result = ranker.run(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
|
||||
def test_init_invalid_custom_prompt_raises(mock_chat_generator):
|
||||
with pytest.raises(TemplateSyntaxError):
|
||||
LLMRanker(chat_generator=mock_chat_generator, prompt="Rank {{ query }")
|
||||
|
||||
|
||||
def test_init_prompt_requires_query_and_documents(mock_chat_generator):
|
||||
with pytest.raises(ValueError, match="prompt must include exactly the variables 'documents' and 'query'"):
|
||||
LLMRanker(chat_generator=mock_chat_generator, prompt="Rank {{ query }}")
|
||||
|
||||
|
||||
def test_init_prompt_rejects_additional_variables(mock_chat_generator):
|
||||
with pytest.raises(ValueError, match="prompt must include exactly the variables 'documents' and 'query'"):
|
||||
LLMRanker(
|
||||
chat_generator=mock_chat_generator,
|
||||
prompt="Rank {{ query }} using {{ documents|length }} docs with top_k={{ top_k }}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY", None),
|
||||
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
|
||||
)
|
||||
def test_live_run_ranks_berlin_first_for_germany_query():
|
||||
documents = [
|
||||
Document(id="doc-berlin", content="Berlin is the capital of Germany."),
|
||||
Document(id="doc-paris", content="Paris is the capital of France."),
|
||||
Document(id="doc-rust", content="Rust is a systems programming language focused on safety."),
|
||||
]
|
||||
ranker = LLMRanker(top_k=2)
|
||||
|
||||
result = ranker.run(query="What is the capital of Germany?", documents=documents)
|
||||
|
||||
assert result["documents"]
|
||||
assert result["documents"][0].id == "doc-berlin"
|
||||
assert len(result["documents"]) <= 2
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY", None),
|
||||
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
|
||||
)
|
||||
def test_live_run_ranks_rust_for_programming_language_query():
|
||||
documents = [
|
||||
Document(id="doc-berlin", content="Berlin is the capital of Germany."),
|
||||
Document(id="doc-paris", content="Paris is the capital of France."),
|
||||
Document(id="doc-rust", content="Rust is a systems programming language focused on safety."),
|
||||
]
|
||||
ranker = LLMRanker(top_k=1)
|
||||
|
||||
result = ranker.run(query="Which document is about a programming language?", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["doc-rust"]
|
||||
|
||||
|
||||
class FakeSyncOnlyChatGenerator:
|
||||
"""A chat generator exposing only a synchronous `run` (no `run_async`) for the fallback path."""
|
||||
|
||||
def __init__(self):
|
||||
self.run = Mock()
|
||||
|
||||
|
||||
class TestLLMRankerAsync:
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(self):
|
||||
documents = [
|
||||
Document(id="1", content="first"),
|
||||
Document(id="2", content="second"),
|
||||
Document(id="3", content="third"),
|
||||
]
|
||||
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
|
||||
mock_chat_generator.run_async = AsyncMock(
|
||||
return_value={
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')]
|
||||
}
|
||||
)
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
|
||||
|
||||
result = await ranker.run_async(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2", "1"]
|
||||
mock_chat_generator.run_async.assert_awaited_once()
|
||||
mock_chat_generator.run.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_fallback_to_sync_run(self):
|
||||
documents = [
|
||||
Document(id="1", content="first"),
|
||||
Document(id="2", content="second"),
|
||||
Document(id="3", content="third"),
|
||||
]
|
||||
fake_chat_generator = FakeSyncOnlyChatGenerator()
|
||||
fake_chat_generator.run.return_value = {
|
||||
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')]
|
||||
}
|
||||
assert not hasattr(fake_chat_generator, "run_async")
|
||||
ranker = LLMRanker(chat_generator=fake_chat_generator, top_k=2)
|
||||
|
||||
result = await ranker.run_async(query="test query", documents=documents)
|
||||
|
||||
assert [document.id for document in result["documents"]] == ["2", "1"]
|
||||
fake_chat_generator.run.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_generator_exception_falls_back(self):
|
||||
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
|
||||
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
|
||||
mock_chat_generator.run_async = AsyncMock(side_effect=RuntimeError("generator failed"))
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1, raise_on_failure=False)
|
||||
|
||||
result = await ranker.run_async(query="test query", documents=documents)
|
||||
|
||||
assert result == {"documents": documents}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_generator_exception_raises(self):
|
||||
documents = [Document(id="1", content="first")]
|
||||
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
|
||||
mock_chat_generator.run_async = AsyncMock(side_effect=RuntimeError("generator failed"))
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator, raise_on_failure=True)
|
||||
|
||||
with pytest.raises(RuntimeError, match="generator failed"):
|
||||
await ranker.run_async(query="test query", documents=documents)
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY", None),
|
||||
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_async_ranks_berlin_first_for_germany_query(self):
|
||||
documents = [
|
||||
Document(id="doc-berlin", content="Berlin is the capital of Germany."),
|
||||
Document(id="doc-paris", content="Paris is the capital of France."),
|
||||
Document(id="doc-rust", content="Rust is a systems programming language focused on safety."),
|
||||
]
|
||||
ranker = LLMRanker(top_k=2)
|
||||
|
||||
result = await ranker.run_async(query="What is the capital of Germany?", documents=documents)
|
||||
|
||||
assert result["documents"]
|
||||
assert result["documents"][0].id == "doc-berlin"
|
||||
assert len(result["documents"]) <= 2
|
||||
|
||||
|
||||
class TestComponentLifecycle:
|
||||
def test_warm_up_delegates_to_chat_generator(self, mock_chat_generator):
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
ranker.warm_up()
|
||||
mock_chat_generator.warm_up.assert_called_once()
|
||||
|
||||
async def test_warm_up_async_delegates_to_chat_generator(self, mock_chat_generator):
|
||||
mock_chat_generator.warm_up_async = AsyncMock()
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
await ranker.warm_up_async()
|
||||
mock_chat_generator.warm_up_async.assert_awaited_once()
|
||||
|
||||
async def test_warm_up_async_falls_back_to_sync_warm_up(self):
|
||||
chat_generator = Mock(spec=["run", "warm_up"])
|
||||
ranker = LLMRanker(chat_generator=chat_generator)
|
||||
await ranker.warm_up_async()
|
||||
chat_generator.warm_up.assert_called_once()
|
||||
|
||||
def test_close_delegates_to_chat_generator(self, mock_chat_generator):
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
ranker.close()
|
||||
mock_chat_generator.close.assert_called_once()
|
||||
|
||||
async def test_close_async_delegates_to_chat_generator(self, mock_chat_generator):
|
||||
mock_chat_generator.close_async = AsyncMock()
|
||||
ranker = LLMRanker(chat_generator=mock_chat_generator)
|
||||
await ranker.close_async()
|
||||
mock_chat_generator.close_async.assert_awaited_once()
|
||||
|
||||
async def test_close_async_falls_back_to_sync_close(self):
|
||||
chat_generator = Mock(spec=["run", "close"])
|
||||
ranker = LLMRanker(chat_generator=chat_generator)
|
||||
await ranker.close_async()
|
||||
chat_generator.close.assert_called_once()
|
||||
|
||||
def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self):
|
||||
chat_generator = Mock(spec=["run"])
|
||||
ranker = LLMRanker(chat_generator=chat_generator)
|
||||
ranker.warm_up()
|
||||
ranker.close()
|
||||
@@ -0,0 +1,114 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document
|
||||
from haystack.components.rankers.lost_in_the_middle import LostInTheMiddleRanker
|
||||
|
||||
|
||||
class TestLostInTheMiddleRanker:
|
||||
def test_lost_in_the_middle_order_odd(self):
|
||||
# tests that lost_in_the_middle order works with an odd number of documents
|
||||
docs = [Document(content=str(i)) for i in range(1, 10)]
|
||||
ranker = LostInTheMiddleRanker()
|
||||
result = ranker.run(documents=docs)
|
||||
assert result["documents"]
|
||||
expected_order = ["1", "3", "5", "7", "9", "8", "6", "4", "2"]
|
||||
assert all(doc.content == expected_order[idx] for idx, doc in enumerate(result["documents"]))
|
||||
|
||||
def test_lost_in_the_middle_order_even(self):
|
||||
# tests that lost_in_the_middle order works with an even number of documents
|
||||
docs = [Document(content=str(i)) for i in range(1, 11)]
|
||||
ranker = LostInTheMiddleRanker()
|
||||
result = ranker.run(documents=docs)
|
||||
expected_order = ["1", "3", "5", "7", "9", "10", "8", "6", "4", "2"]
|
||||
assert all(doc.content == expected_order[idx] for idx, doc in enumerate(result["documents"]))
|
||||
|
||||
def test_lost_in_the_middle_order_two_docs(self):
|
||||
# tests that lost_in_the_middle order works with two documents
|
||||
ranker = LostInTheMiddleRanker()
|
||||
# two docs
|
||||
docs = [Document(content="1"), Document(content="2")]
|
||||
result = ranker.run(documents=docs)
|
||||
assert result["documents"][0].content == "1"
|
||||
assert result["documents"][1].content == "2"
|
||||
|
||||
def test_lost_in_the_middle_init(self):
|
||||
# tests that LostInTheMiddleRanker initializes with default values
|
||||
ranker = LostInTheMiddleRanker()
|
||||
assert ranker.word_count_threshold is None
|
||||
|
||||
ranker = LostInTheMiddleRanker(word_count_threshold=10)
|
||||
assert ranker.word_count_threshold == 10
|
||||
|
||||
def test_lost_in_the_middle_init_invalid_word_count_threshold(self):
|
||||
# tests that LostInTheMiddleRanker raises an error when word_count_threshold is <= 0
|
||||
with pytest.raises(ValueError, match="Invalid value for word_count_threshold"):
|
||||
LostInTheMiddleRanker(word_count_threshold=0)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid value for word_count_threshold"):
|
||||
LostInTheMiddleRanker(word_count_threshold=-5)
|
||||
|
||||
def test_lost_in_the_middle_with_word_count_threshold(self):
|
||||
# tests that lost_in_the_middle with word_count_threshold works as expected
|
||||
ranker = LostInTheMiddleRanker(word_count_threshold=6)
|
||||
docs = [Document(content="word" + str(i)) for i in range(1, 10)]
|
||||
# result, _ = ranker.run(query="", documents=docs)
|
||||
result = ranker.run(documents=docs)
|
||||
expected_order = ["word1", "word3", "word5", "word6", "word4", "word2"]
|
||||
assert all(doc.content == expected_order[idx] for idx, doc in enumerate(result["documents"]))
|
||||
|
||||
ranker = LostInTheMiddleRanker(word_count_threshold=9)
|
||||
# result, _ = ranker.run(query="", documents=docs)
|
||||
result = ranker.run(documents=docs)
|
||||
expected_order = ["word1", "word3", "word5", "word7", "word9", "word8", "word6", "word4", "word2"]
|
||||
assert all(doc.content == expected_order[idx] for idx, doc in enumerate(result["documents"]))
|
||||
|
||||
def test_word_count_threshold_greater_than_total_number_of_words_returns_all_documents(self):
|
||||
ranker = LostInTheMiddleRanker(word_count_threshold=100)
|
||||
docs = [Document(content="word" + str(i)) for i in range(1, 10)]
|
||||
ordered_docs = ranker.run(documents=docs)
|
||||
# assert len(ordered_docs) == len(docs)
|
||||
expected_order = ["word1", "word3", "word5", "word7", "word9", "word8", "word6", "word4", "word2"]
|
||||
assert all(doc.content == expected_order[idx] for idx, doc in enumerate(ordered_docs["documents"]))
|
||||
|
||||
def test_empty_documents_returns_empty_list(self):
|
||||
ranker = LostInTheMiddleRanker()
|
||||
result = ranker.run(documents=[])
|
||||
assert result == {"documents": []}
|
||||
|
||||
def test_list_of_one_document_returns_same_document(self):
|
||||
ranker = LostInTheMiddleRanker()
|
||||
doc = Document(content="test")
|
||||
assert ranker.run(documents=[doc]) == {"documents": [doc]}
|
||||
|
||||
def test_run_deduplicates_documents(self):
|
||||
ranker = LostInTheMiddleRanker()
|
||||
docs = [
|
||||
Document(id="duplicate", content="keep me", score=0.9),
|
||||
Document(id="duplicate", content="drop me", score=0.1),
|
||||
Document(id="unique", content="unique"),
|
||||
]
|
||||
result = ranker.run(documents=docs)
|
||||
|
||||
assert len(result["documents"]) == 2
|
||||
assert result["documents"][0].content == "keep me"
|
||||
assert result["documents"][1].content == "unique"
|
||||
|
||||
@pytest.mark.parametrize("top_k", [1, 2, 3, 4, 5, 6, 7, 8, 12, 20])
|
||||
def test_lost_in_the_middle_order_with_top_k(self, top_k: int):
|
||||
# tests that lost_in_the_middle order works with an odd number of documents and a top_k parameter
|
||||
docs = [Document(content=str(i)) for i in range(1, 10)]
|
||||
ranker = LostInTheMiddleRanker()
|
||||
result = ranker.run(documents=docs, top_k=top_k)
|
||||
if top_k < len(docs):
|
||||
# top_k is less than the number of documents, so only the top_k documents should be returned in LITM order
|
||||
assert len(result["documents"]) == top_k
|
||||
expected_order = ranker.run(documents=[Document(content=str(i)) for i in range(1, top_k + 1)])
|
||||
assert result == expected_order
|
||||
else:
|
||||
# top_k is greater than the number of documents, so all documents should be returned in LITM order
|
||||
assert len(result["documents"]) == len(docs)
|
||||
assert result == ranker.run(documents=docs)
|
||||
@@ -0,0 +1,218 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
|
||||
from haystack import Pipeline
|
||||
from haystack.components.rankers.meta_field_grouping_ranker import MetaFieldGroupingRanker
|
||||
from haystack.dataclasses import Document
|
||||
|
||||
DOC_LIST = [
|
||||
# regular
|
||||
Document(content="Javascript is a popular language", meta={"group": "42", "split_id": 7, "subgroup": "subB"}),
|
||||
Document(content="A chromosome is a package of DNA", meta={"group": "314", "split_id": 2, "subgroup": "subC"}),
|
||||
Document(content="DNA carries genetic information", meta={"group": "314", "split_id": 1, "subgroup": "subE"}),
|
||||
Document(content="Blue whales have a big heart", meta={"group": "11", "split_id": 8, "subgroup": "subF"}),
|
||||
Document(content="Python is a popular language", meta={"group": "42", "split_id": 4, "subgroup": "subB"}),
|
||||
Document(content="bla bla bla bla", meta={"split_id": 8, "subgroup": "subG"}),
|
||||
Document(content="Java is a popular programming language", meta={"group": "42", "split_id": 3, "subgroup": "subB"}),
|
||||
Document(content="An octopus has three hearts", meta={"group": "11", "split_id": 2, "subgroup": "subD"}),
|
||||
# without split id
|
||||
Document(content="without split id", meta={"group": "11"}),
|
||||
Document(content="without split id2", meta={"group": "22", "subgroup": "subI"}),
|
||||
Document(content="without split id3", meta={"group": "11"}),
|
||||
# with list values in the metadata
|
||||
Document(content="list values", meta={"value_list": ["11"], "split_id": 8, "sub_value_list": ["subF"]}),
|
||||
Document(content="list values2", meta={"value_list": ["12"], "split_id": 3, "sub_value_list": ["subX"]}),
|
||||
Document(content="list values3", meta={"value_list": ["12"], "split_id": 8, "sub_value_list": ["subX"]}),
|
||||
]
|
||||
|
||||
|
||||
class TestMetaFieldGroupingRanker:
|
||||
def test_init_default(self) -> None:
|
||||
"""
|
||||
Test the default initialization of the MetaFieldGroupingRanker component.
|
||||
"""
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group", sort_docs_by=None)
|
||||
result = sample_ranker.run(documents=[])
|
||||
assert "documents" in result
|
||||
assert result["documents"] == []
|
||||
|
||||
def test_run_group_by_only(self) -> None:
|
||||
"""
|
||||
Test the MetaFieldGroupingRanker component with only the 'group_by' parameter. No subgroup or sorting is done.
|
||||
"""
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group")
|
||||
result = sample_ranker.run(documents=DOC_LIST)
|
||||
assert "documents" in result
|
||||
assert len(DOC_LIST) == len(result["documents"])
|
||||
assert result["documents"][0].meta["split_id"] == 7 and result["documents"][0].meta["group"] == "42"
|
||||
assert result["documents"][1].meta["split_id"] == 4 and result["documents"][1].meta["group"] == "42"
|
||||
assert result["documents"][2].meta["split_id"] == 3 and result["documents"][2].meta["group"] == "42"
|
||||
assert result["documents"][3].meta["split_id"] == 2 and result["documents"][3].meta["group"] == "314"
|
||||
assert result["documents"][4].meta["split_id"] == 1 and result["documents"][4].meta["group"] == "314"
|
||||
assert result["documents"][5].meta["split_id"] == 8 and result["documents"][5].meta["group"] == "11"
|
||||
assert result["documents"][6].meta["split_id"] == 2 and result["documents"][6].meta["group"] == "11"
|
||||
assert result["documents"][7].content == "without split id" and result["documents"][7].meta["group"] == "11"
|
||||
assert result["documents"][8].content == "without split id3" and result["documents"][8].meta["group"] == "11"
|
||||
assert result["documents"][9].content == "without split id2" and result["documents"][9].meta["group"] == "22"
|
||||
assert result["documents"][10].content == "bla bla bla bla"
|
||||
|
||||
def test_with_group_subgroup_and_sorting(self) -> None:
|
||||
"""
|
||||
Test the MetaFieldGroupingRanker component with all parameters set, i.e.: grouping by 'group', subgrouping by
|
||||
'subgroup', and sorting by 'split_id'.
|
||||
"""
|
||||
ranker = MetaFieldGroupingRanker(group_by="group", subgroup_by="subgroup", sort_docs_by="split_id")
|
||||
result = ranker.run(documents=DOC_LIST)
|
||||
|
||||
assert "documents" in result
|
||||
assert len(DOC_LIST) == len(result["documents"])
|
||||
assert (
|
||||
result["documents"][0].meta["subgroup"] == "subB"
|
||||
and result["documents"][0].meta["group"] == "42"
|
||||
and result["documents"][0].meta["split_id"] == 3
|
||||
)
|
||||
assert (
|
||||
result["documents"][1].meta["subgroup"] == "subB"
|
||||
and result["documents"][1].meta["group"] == "42"
|
||||
and result["documents"][1].meta["split_id"] == 4
|
||||
)
|
||||
assert (
|
||||
result["documents"][2].meta["subgroup"] == "subB"
|
||||
and result["documents"][2].meta["group"] == "42"
|
||||
and result["documents"][2].meta["split_id"] == 7
|
||||
)
|
||||
assert result["documents"][3].meta["subgroup"] == "subC" and result["documents"][3].meta["group"] == "314"
|
||||
assert result["documents"][4].meta["subgroup"] == "subE" and result["documents"][4].meta["group"] == "314"
|
||||
assert result["documents"][5].meta["subgroup"] == "subF" and result["documents"][6].meta["group"] == "11"
|
||||
assert result["documents"][6].meta["subgroup"] == "subD" and result["documents"][5].meta["group"] == "11"
|
||||
assert result["documents"][7].content == "without split id" and result["documents"][7].meta["group"] == "11"
|
||||
assert result["documents"][8].content == "without split id3" and result["documents"][8].meta["group"] == "11"
|
||||
assert result["documents"][9].content == "without split id2" and result["documents"][9].meta["group"] == "22"
|
||||
assert result["documents"][10].content == "bla bla bla bla"
|
||||
|
||||
def test_run_with_lists(self) -> None:
|
||||
"""
|
||||
Test if the MetaFieldGroupingRanker component can handle list values in the metadata.
|
||||
"""
|
||||
ranker = MetaFieldGroupingRanker(group_by="value_list", subgroup_by="sub_value_list", sort_docs_by="split_id")
|
||||
result = ranker.run(documents=DOC_LIST)
|
||||
assert "documents" in result
|
||||
assert len(DOC_LIST) == len(result["documents"])
|
||||
assert result["documents"][0].content == "list values" and result["documents"][0].meta["value_list"] == ["11"]
|
||||
assert result["documents"][1].content == "list values2" and result["documents"][1].meta["value_list"] == ["12"]
|
||||
assert result["documents"][2].content == "list values3" and result["documents"][2].meta["value_list"] == ["12"]
|
||||
|
||||
def test_run_empty_input(self) -> None:
|
||||
"""
|
||||
Test the behavior of the MetaFieldGroupingRanker component with an empty list of documents.
|
||||
"""
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group")
|
||||
result = sample_ranker.run(documents=[])
|
||||
assert "documents" in result
|
||||
assert result["documents"] == []
|
||||
|
||||
def test_run_missing_metadata_keys(self) -> None:
|
||||
"""
|
||||
Test the behavior of the MetaFieldGroupingRanker component when some documents are missing the required
|
||||
metadata keys.
|
||||
"""
|
||||
docs_with_missing_keys = [
|
||||
Document(content="Document without group", meta={"split_id": 1, "subgroup": "subA"}),
|
||||
Document(content="Document without subgroup", meta={"group": "42", "split_id": 2}),
|
||||
Document(content="Document with all keys", meta={"group": "42", "split_id": 3, "subgroup": "subB"}),
|
||||
]
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group", subgroup_by="subgroup", sort_docs_by="split_id")
|
||||
result = sample_ranker.run(documents=docs_with_missing_keys)
|
||||
assert "documents" in result
|
||||
assert len(result["documents"]) == 3
|
||||
assert result["documents"][0].meta["group"] == "42"
|
||||
assert result["documents"][1].meta["group"] == "42"
|
||||
assert result["documents"][2].content == "Document without group"
|
||||
|
||||
def test_run_sort_docs_by_non_numeric_field_with_missing_values(self) -> None:
|
||||
"""
|
||||
Test that sorting by a non-numeric metadata field does not raise an error when some documents are missing
|
||||
that field. Documents missing the sort field are placed at the end of their group.
|
||||
"""
|
||||
docs = [
|
||||
Document(content="newest", meta={"group": "42", "date": "2023-03-01"}),
|
||||
Document(content="missing date", meta={"group": "42"}),
|
||||
Document(content="oldest", meta={"group": "42", "date": "2023-01-01"}),
|
||||
]
|
||||
ranker = MetaFieldGroupingRanker(group_by="group", sort_docs_by="date")
|
||||
result = ranker.run(documents=docs)
|
||||
assert "documents" in result
|
||||
assert len(result["documents"]) == 3
|
||||
assert result["documents"][0].content == "oldest"
|
||||
assert result["documents"][1].content == "newest"
|
||||
assert result["documents"][2].content == "missing date"
|
||||
|
||||
def test_run_sort_docs_by_field_present_but_none(self) -> None:
|
||||
"""
|
||||
Test that sorting by a metadata field works when the field is present but set to None for some documents.
|
||||
Documents with a None value are treated like missing values and placed at the end of their group.
|
||||
"""
|
||||
docs = [
|
||||
Document(content="present", meta={"group": "42", "date": "2023-01-01"}),
|
||||
Document(content="none value", meta={"group": "42", "date": None}),
|
||||
Document(content="missing", meta={"group": "42"}),
|
||||
]
|
||||
ranker = MetaFieldGroupingRanker(group_by="group", sort_docs_by="date")
|
||||
result = ranker.run(documents=docs)
|
||||
assert "documents" in result
|
||||
assert len(result["documents"]) == 3
|
||||
assert result["documents"][0].content == "present"
|
||||
assert result["documents"][1].content == "none value"
|
||||
assert result["documents"][2].content == "missing"
|
||||
|
||||
def test_run_metadata_with_different_data_types(self) -> None:
|
||||
"""
|
||||
Test the behavior of the MetaFieldGroupingRanker component when the metadata values have different data types.
|
||||
"""
|
||||
docs_with_mixed_data_types = [
|
||||
Document(content="Document with string group", meta={"group": "42", "split_id": 1, "subgroup": "subA"}),
|
||||
Document(content="Document with number group", meta={"group": 42, "split_id": 2, "subgroup": "subB"}),
|
||||
Document(content="Document with boolean group", meta={"group": True, "split_id": 3, "subgroup": "subC"}),
|
||||
]
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group", subgroup_by="subgroup", sort_docs_by="split_id")
|
||||
result = sample_ranker.run(documents=docs_with_mixed_data_types)
|
||||
assert "documents" in result
|
||||
assert len(result["documents"]) == 3
|
||||
assert result["documents"][0].meta["group"] == "42"
|
||||
assert result["documents"][1].meta["group"] == 42
|
||||
assert result["documents"][2].meta["group"] is True
|
||||
|
||||
def test_run_deduplicates_documents(self) -> None:
|
||||
"""
|
||||
Test that duplicate documents are removed before grouping.
|
||||
"""
|
||||
docs_with_duplicates = [
|
||||
Document(id="duplicate", content="keep me", meta={"group": "42", "split_id": 1, "subgroup": "subA"}),
|
||||
Document(id="duplicate", content="drop me", meta={"group": "42", "split_id": 1, "subgroup": "subA"}),
|
||||
Document(id="unique", content="unique", meta={"group": "42", "split_id": 2, "subgroup": "subB"}),
|
||||
Document(id="unique2", content="unique2", meta={"group": "42", "split_id": 2, "subgroup": "subA"}),
|
||||
]
|
||||
sample_ranker = MetaFieldGroupingRanker(group_by="group", subgroup_by="subgroup", sort_docs_by="split_id")
|
||||
result = sample_ranker.run(documents=docs_with_duplicates)
|
||||
assert "documents" in result
|
||||
assert len(result["documents"]) == 3
|
||||
assert result["documents"][0].content == "keep me"
|
||||
assert result["documents"][1].content == "unique2"
|
||||
assert result["documents"][2].content == "unique"
|
||||
|
||||
def test_run_in_pipeline_dumps_and_loads(self) -> None:
|
||||
"""
|
||||
Test if the MetaFieldGroupingRanker component can be dumped to a YAML string and reloaded from it.
|
||||
"""
|
||||
ranker = MetaFieldGroupingRanker(group_by="group", sort_docs_by="split_id")
|
||||
result_single = ranker.run(documents=DOC_LIST)
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("ranker", ranker)
|
||||
pipeline_yaml_str = pipeline.dumps()
|
||||
pipeline_reloaded = Pipeline().loads(pipeline_yaml_str)
|
||||
result: dict[str, Any] = pipeline_reloaded.run(data={"documents": DOC_LIST})
|
||||
result = result["ranker"]
|
||||
assert result_single == result
|
||||
@@ -0,0 +1,306 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document
|
||||
from haystack.components.rankers.meta_field import MetaFieldRanker
|
||||
|
||||
|
||||
class TestMetaFieldRanker:
|
||||
@pytest.mark.parametrize("meta_field_values, expected_first_value", [([1.3, 0.7, 2.1], 2.1), ([1, 5, 8], 8)])
|
||||
def test_run(self, meta_field_values, expected_first_value):
|
||||
"""
|
||||
Test if the component ranks documents correctly.
|
||||
"""
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in meta_field_values]
|
||||
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
|
||||
assert len(docs_after) == 3
|
||||
assert docs_after[0].meta["rating"] == expected_first_value
|
||||
|
||||
sorted_scores = sorted([doc.meta["rating"] for doc in docs_after], reverse=True)
|
||||
assert [doc.meta["rating"] for doc in docs_after] == sorted_scores
|
||||
|
||||
def test_run_with_weight_equal_to_0(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=0.0)
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in [1.1, 0.5, 2.3]]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
|
||||
assert len(docs_after) == 3
|
||||
assert [doc.meta["rating"] for doc in docs_after] == [1.1, 0.5, 2.3]
|
||||
|
||||
def test_run_with_weight_equal_to_1(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=1.0)
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in [1.1, 0.5, 2.3]]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
|
||||
assert len(docs_after) == 3
|
||||
sorted_scores = sorted([doc.meta["rating"] for doc in docs_after], reverse=True)
|
||||
assert [doc.meta["rating"] for doc in docs_after] == sorted_scores
|
||||
|
||||
def test_run_with_weight_equal_to_1_passed_in_run_method(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=0.0)
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in [1.1, 0.5, 2.3]]
|
||||
output = ranker.run(documents=docs_before, weight=1.0)
|
||||
docs_after = output["documents"]
|
||||
|
||||
assert len(docs_after) == 3
|
||||
sorted_scores = sorted([doc.meta["rating"] for doc in docs_after], reverse=True)
|
||||
assert [doc.meta["rating"] for doc in docs_after] == sorted_scores
|
||||
|
||||
def test_sort_order_ascending(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=1.0, sort_order="ascending")
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in [1.1, 0.5, 2.3]]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
|
||||
assert len(docs_after) == 3
|
||||
sorted_scores = sorted([doc.meta["rating"] for doc in docs_after])
|
||||
assert [doc.meta["rating"] for doc in docs_after] == sorted_scores
|
||||
|
||||
def test_meta_value_type_float(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=1.0, meta_value_type="float")
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in ["1.1", "10.5", "2.3"]]
|
||||
docs_after = ranker.run(documents=docs_before)["documents"]
|
||||
assert len(docs_after) == 3
|
||||
assert [doc.meta["rating"] for doc in docs_after] == ["10.5", "2.3", "1.1"]
|
||||
|
||||
def test_run_deduplicates_documents(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
docs_before = [
|
||||
Document(id="duplicate", content="keep me", meta={"rating": 1.3}, score=0.9),
|
||||
Document(id="duplicate", content="drop me", meta={"rating": 1.2}, score=0.1),
|
||||
Document(id="unique", content="unique", meta={"rating": 2.1}),
|
||||
]
|
||||
result = ranker.run(documents=docs_before)
|
||||
assert len(result["documents"]) == 2
|
||||
assert result["documents"][0].content == "unique"
|
||||
assert result["documents"][1].content == "keep me"
|
||||
|
||||
def test_run_deduplicates_documents_when_weight_is_0(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=0)
|
||||
docs_before = [
|
||||
Document(id="duplicate", content="keep me", meta={"rating": 1.3}, score=0.9),
|
||||
Document(id="duplicate", content="drop me", meta={"rating": 1.2}, score=0.1),
|
||||
Document(id="unique", content="unique", meta={"rating": 2.1}),
|
||||
]
|
||||
result = ranker.run(documents=docs_before)
|
||||
assert len(result["documents"]) == 2
|
||||
assert result["documents"][0].content == "keep me"
|
||||
assert result["documents"][1].content == "unique"
|
||||
|
||||
def test_meta_value_type_int(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=1.0, meta_value_type="int")
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in ["1", "10", "2"]]
|
||||
docs_after = ranker.run(documents=docs_before)["documents"]
|
||||
assert len(docs_after) == 3
|
||||
assert [doc.meta["rating"] for doc in docs_after] == ["10", "2", "1"]
|
||||
|
||||
def test_meta_value_type_date(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", weight=1.0, meta_value_type="date")
|
||||
docs_before = [Document(content="abc", meta={"rating": value}) for value in ["2022-10", "2023-01", "2022-11"]]
|
||||
docs_after = ranker.run(documents=docs_before)["documents"]
|
||||
assert len(docs_after) == 3
|
||||
assert [doc.meta["rating"] for doc in docs_after] == ["2023-01", "2022-11", "2022-10"]
|
||||
|
||||
def test_returns_empty_list_if_no_documents_are_provided(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
output = ranker.run(documents=[])
|
||||
docs_after = output["documents"]
|
||||
assert docs_after == []
|
||||
|
||||
def test_warning_if_meta_not_found(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
docs_before = [Document(id="1", content="abc", meta={"wrong_field": 1.3})]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
ranker.run(documents=docs_before)
|
||||
assert (
|
||||
"The parameter <meta_field> is currently set to 'rating', but none of the provided Documents with IDs "
|
||||
"1 have this meta key." in caplog.text
|
||||
)
|
||||
|
||||
def test_warning_if_some_meta_not_found(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"wrong_field": 1.3}),
|
||||
Document(id="2", content="def", meta={"rating": 1.3}),
|
||||
]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
ranker.run(documents=docs_before)
|
||||
assert (
|
||||
"The parameter <meta_field> is currently set to 'rating' but the Documents with IDs 1 don't have "
|
||||
"this meta key." in caplog.text
|
||||
)
|
||||
|
||||
def test_warning_if_unsortable_values(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": 1.3}),
|
||||
Document(id="2", content="abc", meta={"rating": "1.2"}),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}),
|
||||
]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
output = ranker.run(documents=docs_before)
|
||||
assert len(output["documents"]) == 3
|
||||
assert "Tried to sort Documents with IDs 1,2,3, but got TypeError with the message:" in caplog.text
|
||||
|
||||
def test_warning_if_meta_value_parsing_error(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating", meta_value_type="float")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": "1.3"}),
|
||||
Document(id="2", content="abc", meta={"rating": "1.2"}),
|
||||
Document(id="3", content="abc", meta={"rating": "not a float"}),
|
||||
]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
output = ranker.run(documents=docs_before)
|
||||
assert len(output["documents"]) == 3
|
||||
assert (
|
||||
"Tried to parse the meta values of Documents with IDs 1,2,3, but got ValueError with the message:"
|
||||
in caplog.text
|
||||
)
|
||||
|
||||
def test_warning_meta_value_type_not_all_strings(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating", meta_value_type="float")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": "1.3"}),
|
||||
Document(id="2", content="abc", meta={"rating": "1.2"}),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}),
|
||||
]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
output = ranker.run(documents=docs_before)
|
||||
assert len(output["documents"]) == 3
|
||||
assert (
|
||||
"The parameter <meta_value_type> is currently set to 'float', but not all of meta values in the "
|
||||
"provided Documents with IDs 1,2,3 are strings." in caplog.text
|
||||
)
|
||||
|
||||
def test_raises_value_error_if_wrong_ranking_mode(self):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", ranking_mode="wrong_mode")
|
||||
|
||||
def test_raises_value_error_if_wrong_top_k(self):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", top_k=-1)
|
||||
|
||||
@pytest.mark.parametrize("score", [-1, 2, 1.3, 2.1])
|
||||
def test_raises_component_error_if_wrong_weight(self, score):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", weight=score)
|
||||
|
||||
def test_raises_value_error_if_wrong_sort_order(self):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", sort_order="wrong_order")
|
||||
|
||||
def test_raises_value_error_if_wrong_missing_meta(self):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", missing_meta="wrong_missing_meta")
|
||||
|
||||
def test_raises_value_error_if_wrong_meta_value_type(self):
|
||||
with pytest.raises(ValueError):
|
||||
MetaFieldRanker(meta_field="rating", meta_value_type="wrong_type")
|
||||
|
||||
def test_linear_score(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5)
|
||||
docs_before = [
|
||||
Document(content="abc", meta={"rating": 1.3}, score=0.3),
|
||||
Document(content="abc", meta={"rating": 0.7}, score=0.4),
|
||||
Document(content="abc", meta={"rating": 2.1}, score=0.6),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert docs_after[0].score == 0.8
|
||||
|
||||
def test_reciprocal_rank_fusion(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="reciprocal_rank_fusion", weight=0.5)
|
||||
docs_before = [
|
||||
Document(content="abc", meta={"rating": 1.3}, score=0.3),
|
||||
Document(content="abc", meta={"rating": 0.7}, score=0.4),
|
||||
Document(content="abc", meta={"rating": 2.1}, score=0.6),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert docs_after[0].score == pytest.approx(0.016261, abs=1e-5)
|
||||
|
||||
@pytest.mark.parametrize("score", [-1, 2, 1.3, 2.1])
|
||||
def test_linear_score_raises_warning_if_doc_wrong_score(self, score, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5)
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": 1.3}, score=score),
|
||||
Document(id="2", content="abc", meta={"rating": 0.7}, score=0.4),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}, score=0.6),
|
||||
]
|
||||
with caplog.at_level(logging.WARNING):
|
||||
ranker.run(documents=docs_before)
|
||||
assert f"The score {score} for Document 1 is outside the [0,1] range; defaulting to 0" in caplog.text
|
||||
|
||||
def test_linear_score_raises_raises_warning_if_doc_without_score(self, caplog):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5)
|
||||
docs_before = [
|
||||
Document(content="abc", meta={"rating": 1.3}),
|
||||
Document(content="abc", meta={"rating": 0.7}),
|
||||
Document(content="abc", meta={"rating": 2.1}),
|
||||
]
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
ranker.run(documents=docs_before)
|
||||
assert "The score wasn't provided; defaulting to 0." in caplog.text
|
||||
|
||||
def test_different_ranking_mode_for_init_vs_run(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5)
|
||||
docs_before = [
|
||||
Document(content="abc", meta={"rating": 1.3}, score=0.3),
|
||||
Document(content="abc", meta={"rating": 0.7}, score=0.4),
|
||||
Document(content="abc", meta={"rating": 2.1}, score=0.6),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert docs_after[0].score == 0.8
|
||||
|
||||
output = ranker.run(documents=docs_before, ranking_mode="reciprocal_rank_fusion")
|
||||
docs_after = output["documents"]
|
||||
assert docs_after[0].score == pytest.approx(0.016261, abs=1e-5)
|
||||
|
||||
def test_missing_meta_bottom(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5, missing_meta="bottom")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": 1.3}, score=0.6),
|
||||
Document(id="2", content="abc", meta={}, score=0.4),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}, score=0.39),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert len(docs_after) == 3
|
||||
assert docs_after[2].id == "2"
|
||||
|
||||
def test_missing_meta_top(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5, missing_meta="top")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": 1.3}, score=0.6),
|
||||
Document(id="2", content="abc", meta={}, score=0.59),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}, score=0.4),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert len(docs_after) == 3
|
||||
assert docs_after[0].id == "2"
|
||||
|
||||
def test_missing_meta_drop(self):
|
||||
ranker = MetaFieldRanker(meta_field="rating", ranking_mode="linear_score", weight=0.5, missing_meta="drop")
|
||||
docs_before = [
|
||||
Document(id="1", content="abc", meta={"rating": 1.3}, score=0.6),
|
||||
Document(id="2", content="abc", meta={}, score=0.59),
|
||||
Document(id="3", content="abc", meta={"rating": 2.1}, score=0.4),
|
||||
]
|
||||
output = ranker.run(documents=docs_before)
|
||||
docs_after = output["documents"]
|
||||
assert len(docs_after) == 2
|
||||
assert "2" not in [doc.id for doc in docs_after]
|
||||
Reference in New Issue
Block a user