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 logging
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document, GeneratedAnswer
|
||||
from haystack.components.builders.answer_builder import AnswerBuilder
|
||||
from haystack.dataclasses.chat_message import ChatMessage
|
||||
|
||||
|
||||
def _check_metadata_excluding_all_messages(actual_meta, expected_meta):
|
||||
"""Helper function to check metadata while ignoring the all_messages field"""
|
||||
actual_without_messages = {k: v for k, v in actual_meta.items() if k != "all_messages"}
|
||||
assert actual_without_messages == expected_meta
|
||||
|
||||
|
||||
class TestAnswerBuilder:
|
||||
def test_run_unmatching_input_len(self):
|
||||
component = AnswerBuilder()
|
||||
with pytest.raises(ValueError):
|
||||
component.run(query="query", replies=["reply1"], meta=[{"test": "meta"}, {"test": "meta2"}])
|
||||
|
||||
def test_run_without_meta(self):
|
||||
component = AnswerBuilder()
|
||||
output = component.run(query="query", replies=["reply1"])
|
||||
answers = output["answers"]
|
||||
assert answers[0].data == "reply1"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta # Check that all_messages exists
|
||||
assert answers[0].query == "query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_meta_is_an_empty_list(self):
|
||||
component = AnswerBuilder()
|
||||
output = component.run(query="query", replies=["reply1"], meta=[])
|
||||
answers = output["answers"]
|
||||
assert answers[0].data == "reply1"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_meta(self):
|
||||
component = AnswerBuilder()
|
||||
output = component.run(query="query", replies=["reply1"], meta=[{"test": "meta"}])
|
||||
answers = output["answers"]
|
||||
assert answers[0].data == "reply1"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {"test": "meta"})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].meta["all_messages"] == ["reply1"]
|
||||
assert answers[0].query == "query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_without_pattern(self):
|
||||
component = AnswerBuilder()
|
||||
output = component.run(query="test query", replies=["Answer: AnswerString"], meta=[{}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_pattern_with_capturing_group(self):
|
||||
component = AnswerBuilder(pattern=r"Answer: (.*)")
|
||||
output = component.run(query="test query", replies=["Answer: AnswerString"], meta=[{}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_pattern_without_capturing_group(self):
|
||||
component = AnswerBuilder(pattern=r"'.*'")
|
||||
output = component.run(query="test query", replies=["Answer: 'AnswerString'"], meta=[{}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "'AnswerString'"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_pattern_with_more_than_one_capturing_group(self):
|
||||
with pytest.raises(ValueError, match="contains multiple capture groups"):
|
||||
AnswerBuilder(pattern=r"Answer: (.*), (.*)")
|
||||
|
||||
def test_run_with_pattern_set_at_runtime(self):
|
||||
component = AnswerBuilder(pattern="unused pattern")
|
||||
output = component.run(query="test query", replies=["Answer: AnswerString"], meta=[{}], pattern=r"Answer: (.*)")
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_documents_without_reference_pattern(self):
|
||||
component = AnswerBuilder()
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=["Answer: AnswerString"],
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2")],
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 2
|
||||
assert answers[0].documents[0].content == "test doc 1"
|
||||
assert answers[0].documents[1].content == "test doc 2"
|
||||
|
||||
def test_run_with_documents_with_reference_pattern(self):
|
||||
component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]")
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=["Answer: AnswerString[2]"],
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2")],
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString[2]"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 1
|
||||
assert answers[0].documents[0].content == "test doc 2"
|
||||
assert answers[0].documents[0].meta["referenced"] is True
|
||||
assert answers[0].documents[0].meta["source_index"] == 2
|
||||
|
||||
def test_run_with_documents_with_reference_pattern_return_all_documents(self):
|
||||
component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]", return_only_referenced_documents=False)
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=["Answer: AnswerString[2]"],
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2")],
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString[2]"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 2
|
||||
assert answers[0].documents[0].content == "test doc 1"
|
||||
assert answers[0].documents[0].meta["referenced"] is False
|
||||
assert answers[0].documents[0].meta["source_index"] == 1
|
||||
assert answers[0].documents[1].content == "test doc 2"
|
||||
assert answers[0].documents[1].meta["referenced"] is True
|
||||
assert answers[0].documents[1].meta["source_index"] == 2
|
||||
|
||||
def test_run_with_documents_with_reference_pattern_and_no_match(self, caplog):
|
||||
component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]")
|
||||
with caplog.at_level(logging.WARNING):
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=["Answer: AnswerString[3]"],
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2")],
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString[3]"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 0
|
||||
assert "Document index '3' referenced in Generator output is out of range." in caplog.text
|
||||
|
||||
def test_run_with_reference_pattern_set_at_runtime(self):
|
||||
component = AnswerBuilder(reference_pattern="unused pattern")
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=["Answer: AnswerString[2][3]"],
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2"), Document(content="test doc 3")],
|
||||
reference_pattern="\\[(\\d+)\\]",
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString[2][3]"
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {})
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 2
|
||||
assert answers[0].documents[0].content == "test doc 2"
|
||||
assert answers[0].documents[1].content == "test doc 3"
|
||||
|
||||
def test_run_with_chat_message_replies_without_pattern(self):
|
||||
component = AnswerBuilder()
|
||||
|
||||
message_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
}
|
||||
replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)]
|
||||
|
||||
output = component.run(query="test query", replies=replies, meta=[{}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString"
|
||||
|
||||
# Check metadata excluding all_messages
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, message_meta)
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_chat_message_replies_with_pattern(self):
|
||||
component = AnswerBuilder(pattern=r"Answer: (.*)")
|
||||
|
||||
message_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
}
|
||||
replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)]
|
||||
|
||||
output = component.run(query="test query", replies=replies, meta=[{}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, message_meta)
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_chat_message_replies_with_documents(self):
|
||||
component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]")
|
||||
message_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
}
|
||||
replies = [ChatMessage.from_assistant("Answer: AnswerString[2]", meta=message_meta)]
|
||||
|
||||
output = component.run(
|
||||
query="test query",
|
||||
replies=replies,
|
||||
meta=[{}],
|
||||
documents=[Document(content="test doc 1"), Document(content="test doc 2")],
|
||||
)
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Answer: AnswerString[2]"
|
||||
|
||||
# Check metadata excluding all_messages
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, message_meta)
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert len(answers[0].documents) == 1
|
||||
assert answers[0].documents[0].content == "test doc 2"
|
||||
assert answers[0].meta["all_messages"] == [
|
||||
ChatMessage.from_assistant("Answer: AnswerString[2]", meta=message_meta)
|
||||
]
|
||||
|
||||
def test_run_with_chat_message_replies_with_pattern_set_at_runtime(self):
|
||||
component = AnswerBuilder(pattern="unused pattern")
|
||||
message_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
}
|
||||
replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)]
|
||||
|
||||
output = component.run(query="test query", replies=replies, meta=[{}], pattern=r"Answer: (.*)")
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
|
||||
# Check metadata excluding all_messages
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, message_meta)
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_chat_message_replies_with_meta_set_at_run_time(self):
|
||||
component = AnswerBuilder()
|
||||
message_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
}
|
||||
replies = [ChatMessage.from_assistant("AnswerString", meta=message_meta)]
|
||||
|
||||
output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}])
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
|
||||
# Check metadata excluding all_messages
|
||||
expected_meta = {
|
||||
"model": "gpt-4o-mini",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185},
|
||||
"test": "meta",
|
||||
}
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, expected_meta)
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_run_with_chat_message_no_meta_with_meta_set_at_run_time(self):
|
||||
component = AnswerBuilder()
|
||||
replies = [ChatMessage.from_assistant("AnswerString")]
|
||||
output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}])
|
||||
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "AnswerString"
|
||||
|
||||
# Check metadata excluding all_messages
|
||||
_check_metadata_excluding_all_messages(answers[0].meta, {"test": "meta"})
|
||||
assert "all_messages" in answers[0].meta
|
||||
|
||||
assert answers[0].query == "test query"
|
||||
assert answers[0].documents == []
|
||||
assert isinstance(answers[0], GeneratedAnswer)
|
||||
|
||||
def test_conversation_history_with_last_message_only_false(self):
|
||||
"""Test behavior with last_message_only=False (default)"""
|
||||
# Test with two messages in replies
|
||||
component = AnswerBuilder(last_message_only=False)
|
||||
replies = [
|
||||
ChatMessage.from_user("What is Haystack?"),
|
||||
ChatMessage.from_assistant("Haystack is a framework for building LLM applications."),
|
||||
]
|
||||
output = component.run(query="test query", replies=replies)
|
||||
|
||||
answers = output["answers"]
|
||||
# Should have one answer for each message in replies
|
||||
assert len(answers) == 2
|
||||
assert answers[0].data == "What is Haystack?"
|
||||
assert answers[1].data == "Haystack is a framework for building LLM applications."
|
||||
|
||||
# Check that each answer contains the full conversation history
|
||||
for answer in answers:
|
||||
assert "all_messages" in answer.meta
|
||||
assert len(answer.meta["all_messages"]) == 2
|
||||
assert answer.meta["all_messages"] == replies
|
||||
|
||||
# Test with one message in replies
|
||||
component = AnswerBuilder(last_message_only=False)
|
||||
replies = [ChatMessage.from_user("What is Haystack?")]
|
||||
output = component.run(query="test query", replies=replies)
|
||||
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "What is Haystack?"
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].meta["all_messages"] == replies
|
||||
assert len(answers[0].meta["all_messages"]) == 1
|
||||
|
||||
def test_conversation_history_with_last_message_only_true(self):
|
||||
"""Test behavior with last_message_only=True"""
|
||||
# Test with two messages in replies
|
||||
component = AnswerBuilder(last_message_only=True)
|
||||
replies = [
|
||||
ChatMessage.from_user("What is Haystack?"),
|
||||
ChatMessage.from_assistant("Haystack is a framework for building LLM applications."),
|
||||
]
|
||||
output = component.run(query="test query", replies=replies)
|
||||
|
||||
answers = output["answers"]
|
||||
# Should only have one answer - the last message
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "Haystack is a framework for building LLM applications."
|
||||
|
||||
# Check that the answer contains the full conversation history
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert len(answers[0].meta["all_messages"]) == 2
|
||||
assert answers[0].meta["all_messages"] == replies
|
||||
|
||||
# Test with one message in replies
|
||||
component = AnswerBuilder(last_message_only=True)
|
||||
replies = [ChatMessage.from_user("What is Haystack?")]
|
||||
output = component.run(query="test query", replies=replies)
|
||||
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0].data == "What is Haystack?"
|
||||
assert "all_messages" in answers[0].meta
|
||||
assert answers[0].meta["all_messages"] == replies
|
||||
assert len(answers[0].meta["all_messages"]) == 1
|
||||
|
||||
def test_run_does_not_mutate_input_documents_meta(self):
|
||||
doc = Document(content="Paris is the capital of France.", meta={"source": "wiki"})
|
||||
component = AnswerBuilder()
|
||||
component.run(query="Capital of France?", replies=["Paris."], documents=[doc])
|
||||
assert doc.meta == {"source": "wiki"}
|
||||
|
||||
def test_run_does_not_mutate_input_documents_meta_with_reference_pattern(self):
|
||||
doc = Document(content="Paris is the capital of France.", meta={"source": "wiki"})
|
||||
component = AnswerBuilder(reference_pattern=r"\[(\d+)\]", return_only_referenced_documents=False)
|
||||
component.run(query="Capital?", replies=["Paris [1]."], documents=[doc])
|
||||
assert doc.meta == {"source": "wiki"}
|
||||
|
||||
def test_run_does_not_mutate_document_with_empty_meta(self):
|
||||
doc = Document(content="Paris is the capital of France.")
|
||||
component = AnswerBuilder()
|
||||
component.run(query="Capital of France?", replies=["Paris."], documents=[doc])
|
||||
assert doc.meta == {}
|
||||
|
||||
def test_run_expands_reference_ranges_when_enabled(self):
|
||||
docs = [Document(content=f"doc {i}") for i in range(1, 11)]
|
||||
component = AnswerBuilder(
|
||||
reference_pattern=r"\[(\d+)\]", expand_reference_ranges=True, return_only_referenced_documents=False
|
||||
)
|
||||
output = component.run(query="test query", replies=["Answer citing sources [6-10]."], documents=docs)
|
||||
answer = output["answers"][0]
|
||||
referenced_docs = [doc for doc in answer.documents if doc.meta["referenced"]]
|
||||
assert [doc.meta["source_index"] for doc in referenced_docs] == [6, 7, 8, 9, 10]
|
||||
|
||||
def test_run_expands_comma_separated_reference_ranges(self):
|
||||
docs = [Document(content=f"doc {i}") for i in range(1, 11)]
|
||||
component = AnswerBuilder(
|
||||
reference_pattern=r"\[(\d+)\]", expand_reference_ranges=True, return_only_referenced_documents=False
|
||||
)
|
||||
output = component.run(query="test query", replies=["Answer citing sources [1-3,7-9]."], documents=docs)
|
||||
answer = output["answers"][0]
|
||||
referenced_docs = [doc for doc in answer.documents if doc.meta["referenced"]]
|
||||
assert [doc.meta["source_index"] for doc in referenced_docs] == [1, 2, 3, 7, 8, 9]
|
||||
|
||||
def test_run_ignores_invalid_reference_ranges(self):
|
||||
docs = [Document(content=f"doc {i}") for i in range(1, 5)]
|
||||
component = AnswerBuilder(
|
||||
reference_pattern=r"\[(\d+)\]", expand_reference_ranges=True, return_only_referenced_documents=True
|
||||
)
|
||||
output = component.run(query="test query", replies=["Answer citing sources [3-1]."], documents=docs)
|
||||
assert output["answers"][0].documents == []
|
||||
|
||||
def test_run_clamps_reference_range_to_number_of_documents(self):
|
||||
docs = [Document(content=f"doc {i}") for i in range(1, 4)]
|
||||
component = AnswerBuilder(
|
||||
reference_pattern=r"\[(\d+)\]", expand_reference_ranges=True, return_only_referenced_documents=True
|
||||
)
|
||||
output = component.run(query="test query", replies=["Answer citing sources [1-100]."], documents=docs)
|
||||
referenced_docs = output["answers"][0].documents
|
||||
assert [doc.meta["source_index"] for doc in referenced_docs] == [1, 2, 3]
|
||||
|
||||
def test_run_does_not_expand_reference_ranges_by_default(self):
|
||||
docs = [Document(content=f"doc {i}") for i in range(1, 11)]
|
||||
component = AnswerBuilder(reference_pattern=r"\[(\d+)\]", return_only_referenced_documents=True)
|
||||
output = component.run(query="test query", replies=["Answer citing sources [6-10]."], documents=docs)
|
||||
assert output["answers"][0].documents == []
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,378 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import arrow
|
||||
import pytest
|
||||
from jinja2 import TemplateSyntaxError
|
||||
|
||||
from haystack import component
|
||||
from haystack.components.builders.prompt_builder import PromptBuilder
|
||||
from haystack.core.pipeline.pipeline import Pipeline
|
||||
from haystack.dataclasses.document import Document
|
||||
|
||||
|
||||
class TestPromptBuilder:
|
||||
def test_init(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
assert builder.template is not None
|
||||
assert builder.required_variables == "*"
|
||||
assert builder._template_string == "This is a {{ variable }}"
|
||||
assert builder._variables is None
|
||||
assert builder._required_variables == "*"
|
||||
|
||||
# we have inputs that contain: template, template_variables + inferred variables
|
||||
inputs = builder.__haystack_input__._sockets_dict
|
||||
assert set(inputs.keys()) == {"template", "template_variables", "variable"}
|
||||
assert inputs["template"].type == str | None
|
||||
assert inputs["template_variables"].type == dict[str, Any] | None
|
||||
assert inputs["variable"].type == Any
|
||||
|
||||
# response is always prompt
|
||||
outputs = builder.__haystack_output__._sockets_dict
|
||||
assert set(outputs.keys()) == {"prompt"}
|
||||
assert outputs["prompt"].type == str
|
||||
|
||||
def test_init_with_required_variables(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}", required_variables=["variable"])
|
||||
assert builder.template is not None
|
||||
assert builder.required_variables == ["variable"]
|
||||
assert builder._template_string == "This is a {{ variable }}"
|
||||
assert builder._variables is None
|
||||
assert builder._required_variables == ["variable"]
|
||||
|
||||
# we have inputs that contain: template, template_variables + inferred variables
|
||||
inputs = builder.__haystack_input__._sockets_dict
|
||||
assert set(inputs.keys()) == {"template", "template_variables", "variable"}
|
||||
assert inputs["template"].type == str | None
|
||||
assert inputs["template_variables"].type == dict[str, Any] | None
|
||||
assert inputs["variable"].type == Any
|
||||
|
||||
# response is always prompt
|
||||
outputs = builder.__haystack_output__._sockets_dict
|
||||
assert set(outputs.keys()) == {"prompt"}
|
||||
assert outputs["prompt"].type == str
|
||||
|
||||
def test_init_with_custom_variables(self):
|
||||
variables = ["var1", "var2", "var3"]
|
||||
template = "Hello, {{ var1 }}, {{ var2 }}!"
|
||||
builder = PromptBuilder(template=template, variables=variables)
|
||||
assert builder.template is not None
|
||||
assert builder.required_variables == "*"
|
||||
assert builder._variables == variables
|
||||
assert builder._template_string == "Hello, {{ var1 }}, {{ var2 }}!"
|
||||
assert builder._required_variables == "*"
|
||||
|
||||
# we have inputs that contain: template, template_variables + variables
|
||||
inputs = builder.__haystack_input__._sockets_dict
|
||||
assert set(inputs.keys()) == {"template", "template_variables", "var1", "var2", "var3"}
|
||||
assert inputs["template"].type == str | None
|
||||
assert inputs["template_variables"].type == dict[str, Any] | None
|
||||
assert inputs["var1"].type == Any
|
||||
assert inputs["var2"].type == Any
|
||||
assert inputs["var3"].type == Any
|
||||
|
||||
# response is always prompt
|
||||
outputs = builder.__haystack_output__._sockets_dict
|
||||
assert set(outputs.keys()) == {"prompt"}
|
||||
assert outputs["prompt"].type == str
|
||||
|
||||
def test_to_dict(self):
|
||||
builder = PromptBuilder(
|
||||
template="This is a {{ variable }}", variables=["var1", "var2"], required_variables=["var1", "var3"]
|
||||
)
|
||||
res = builder.to_dict()
|
||||
assert res == {
|
||||
"type": "haystack.components.builders.prompt_builder.PromptBuilder",
|
||||
"init_parameters": {
|
||||
"template": "This is a {{ variable }}",
|
||||
"variables": ["var1", "var2"],
|
||||
"required_variables": ["var1", "var3"],
|
||||
},
|
||||
}
|
||||
|
||||
def test_to_dict_without_optional_params(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
res = builder.to_dict()
|
||||
assert res == {
|
||||
"type": "haystack.components.builders.prompt_builder.PromptBuilder",
|
||||
"init_parameters": {"template": "This is a {{ variable }}", "variables": None, "required_variables": "*"},
|
||||
}
|
||||
|
||||
def test_run(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
res = builder.run(variable="test")
|
||||
assert res == {"prompt": "This is a test"}
|
||||
|
||||
def test_run_template_variable(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
res = builder.run(template_variables={"variable": "test"})
|
||||
assert res == {"prompt": "This is a test"}
|
||||
|
||||
def test_run_template_variable_overrides_variable(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
res = builder.run(template_variables={"variable": "test_from_template_var"}, variable="test")
|
||||
assert res == {"prompt": "This is a test_from_template_var"}
|
||||
|
||||
def test_run_without_input(self):
|
||||
builder = PromptBuilder(template="This is a template without input")
|
||||
res = builder.run()
|
||||
assert res == {"prompt": "This is a template without input"}
|
||||
|
||||
def test_run_with_missing_optional_input(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}", required_variables=None)
|
||||
res = builder.run()
|
||||
assert res == {"prompt": "This is a "}
|
||||
|
||||
def test_run_with_missing_required_input_default(self):
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
with pytest.raises(ValueError, match="variable"):
|
||||
builder.run()
|
||||
|
||||
def test_run_with_missing_required_input(self):
|
||||
builder = PromptBuilder(template="This is a {{ foo }}, not a {{ bar }}", required_variables=["foo", "bar"])
|
||||
with pytest.raises(ValueError, match="foo"):
|
||||
builder.run(bar="bar")
|
||||
with pytest.raises(ValueError, match="bar"):
|
||||
builder.run(foo="foo")
|
||||
with pytest.raises(ValueError, match="foo, bar"):
|
||||
builder.run()
|
||||
|
||||
def test_run_with_missing_required_input_using_star(self):
|
||||
builder = PromptBuilder(template="This is a {{ foo }}, not a {{ bar }}", required_variables="*")
|
||||
with pytest.raises(ValueError, match="foo"):
|
||||
builder.run(bar="bar")
|
||||
with pytest.raises(ValueError, match="bar"):
|
||||
builder.run(foo="foo")
|
||||
with pytest.raises(ValueError, match="bar, foo"):
|
||||
builder.run()
|
||||
|
||||
def test_run_with_variables(self):
|
||||
variables = ["var1", "var2", "var3"]
|
||||
template = "Hello, {{ name }}! {{ var1 }}"
|
||||
builder = PromptBuilder(template=template, variables=variables, required_variables=["name", "var1"])
|
||||
assert builder.run(name="John", var1="How are you?") == {"prompt": "Hello, John! How are you?"}
|
||||
|
||||
def test_run_overwriting_default_template(self):
|
||||
builder = PromptBuilder(template="Hello, {{ name }}!")
|
||||
assert builder.run("Hello, {{ var1 }}{{ name }}!!", name="John") == {"prompt": "Hello, John!!"}
|
||||
|
||||
def test_run_overwriting_default_template_with_template_variables(self):
|
||||
builder = PromptBuilder(template="Hello, {{ name }}!")
|
||||
template = "Hello, {{ var1 }} {{ name }}!"
|
||||
assert builder.run(template, var1="Big", name="John") == {"prompt": "Hello, Big John!"}
|
||||
|
||||
def test_run_overwriting_default_template_with_variables(self):
|
||||
variables = ["var1", "var2", "name"]
|
||||
builder = PromptBuilder(template="Hello, {{ name }}!", variables=variables, required_variables=["name"])
|
||||
template = "Hello, {{ var1 }} {{ name }}!"
|
||||
assert builder.run(template, name="John", var1="Big") == {"prompt": "Hello, Big John!"}
|
||||
|
||||
def test_run_with_invalid_template(self):
|
||||
builder = PromptBuilder(template="Hello, {{ name }}!")
|
||||
|
||||
template = "Hello, {{ name }!"
|
||||
template_variables = {"name": "John"}
|
||||
with pytest.raises(TemplateSyntaxError):
|
||||
builder.run(template, template_variables)
|
||||
|
||||
def test_init_with_invalid_template(self):
|
||||
template = "Hello, {{ name }!"
|
||||
with pytest.raises(TemplateSyntaxError):
|
||||
PromptBuilder(template)
|
||||
|
||||
def test_provided_template_variables(self):
|
||||
prompt_builder = PromptBuilder(template="", variables=["documents"], required_variables=["city"])
|
||||
|
||||
# both variables are provided
|
||||
prompt_builder._validate_variables({"name", "city"})
|
||||
|
||||
# provided variables are a superset of the required variables
|
||||
prompt_builder._validate_variables({"name", "city", "age"})
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
prompt_builder._validate_variables({"name"})
|
||||
|
||||
def test_example_in_pipeline(self):
|
||||
default_template = "Here is the document: {{documents[0].content}} \\n Answer: {{query}}"
|
||||
prompt_builder = PromptBuilder(template=default_template, variables=["documents"])
|
||||
|
||||
@component
|
||||
class DocumentProducer:
|
||||
@component.output_types(documents=list[Document])
|
||||
def run(self, doc_input: str):
|
||||
return {"documents": [Document(content=doc_input)]}
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("doc_producer", DocumentProducer())
|
||||
pipe.add_component("prompt_builder", prompt_builder)
|
||||
pipe.connect("doc_producer.documents", "prompt_builder.documents")
|
||||
|
||||
template = "Here is the document: {{documents[0].content}} \n Query: {{query}}"
|
||||
result = pipe.run(
|
||||
data={
|
||||
"doc_producer": {"doc_input": "Hello world, I live in Berlin"},
|
||||
"prompt_builder": {
|
||||
"template": template,
|
||||
"template_variables": {"query": "Where does the speaker live?"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"prompt_builder": {
|
||||
"prompt": "Here is the document: Hello world, I live in Berlin \n Query: Where does the speaker live?"
|
||||
}
|
||||
}
|
||||
|
||||
def test_example_in_pipeline_simple(self):
|
||||
default_template = "This is the default prompt:\n Query: {{query}}"
|
||||
prompt_builder = PromptBuilder(template=default_template)
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("prompt_builder", prompt_builder)
|
||||
|
||||
# using the default prompt
|
||||
result = pipe.run(data={"query": "Where does the speaker live?"})
|
||||
expected_default = {
|
||||
"prompt_builder": {"prompt": "This is the default prompt:\n Query: Where does the speaker live?"}
|
||||
}
|
||||
assert result == expected_default
|
||||
|
||||
# using the dynamic prompt
|
||||
result = pipe.run(
|
||||
data={"query": "Where does the speaker live?", "template": "This is the dynamic prompt:\n Query: {{query}}"}
|
||||
)
|
||||
expected_dynamic = {
|
||||
"prompt_builder": {"prompt": "This is the dynamic prompt:\n Query: Where does the speaker live?"}
|
||||
}
|
||||
assert result == expected_dynamic
|
||||
|
||||
def test_warning_when_required_variables_explicitly_none(self, caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
_ = PromptBuilder(template="This is a {{ variable }}", required_variables=None)
|
||||
assert "explicitly set to `None`" in caplog.text
|
||||
|
||||
def test_no_warning_with_default_required_variables(self, caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
_ = PromptBuilder(template="This is a {{ variable }}")
|
||||
assert "required_variables" not in caplog.text
|
||||
|
||||
def test_variables_correct_with_assignment(self) -> None:
|
||||
template = """{% if existing_documents is not none %}
|
||||
{% set existing_doc_len = existing_documents|length %}
|
||||
{% else %}
|
||||
{% set existing_doc_len = 0 %}
|
||||
{% endif %}
|
||||
{% for doc in docs %}
|
||||
<document reference="{{loop.index + existing_doc_len}}">
|
||||
{{ doc.content }}
|
||||
</document>
|
||||
{% endfor %}
|
||||
"""
|
||||
builder = PromptBuilder(template=template, required_variables="*")
|
||||
assert set(builder.variables) == {"docs", "existing_documents"}
|
||||
assert builder.required_variables == "*"
|
||||
|
||||
def test_variables_correct_with_tuple_assignment(self):
|
||||
template = """{% if existing_documents is not none -%}
|
||||
{% set x, y = (existing_documents|length, 1) -%}
|
||||
{% else -%}
|
||||
{% set x, y = (0, 1) -%}
|
||||
{% endif -%}
|
||||
x={{ x }}, y={{ y }}
|
||||
"""
|
||||
builder = PromptBuilder(template=template, required_variables="*")
|
||||
assert builder.variables == ["existing_documents"]
|
||||
assert builder.required_variables == "*"
|
||||
res = builder.run(existing_documents=None)
|
||||
assert res["prompt"] == "x=0, y=1"
|
||||
|
||||
def test_variables_correct_with_list_assignment(self):
|
||||
template = """{% if existing_documents is not none -%}
|
||||
{% set x, y = [existing_documents|length, 1] -%}
|
||||
{% else -%}
|
||||
{% set x, y = [0, 1] -%}
|
||||
{% endif -%}
|
||||
x={{ x }}, y={{ y }}
|
||||
"""
|
||||
builder = PromptBuilder(template=template, required_variables="*")
|
||||
assert builder.variables == ["existing_documents"]
|
||||
assert builder.required_variables == "*"
|
||||
res = builder.run(existing_documents=None)
|
||||
assert res["prompt"] == "x=0, y=1"
|
||||
|
||||
|
||||
class TestPromptBuilderWithJinja2TimeExtension:
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_now(self, monkeypatch):
|
||||
"""Mock the arrow.now function to return a fixed datetime"""
|
||||
monkeypatch.setattr("arrow.now", lambda timezone="UTC": arrow.get("1970-01-01 00:00:00").to(timezone))
|
||||
|
||||
@patch("haystack.components.builders.prompt_builder.Jinja2TimeExtension")
|
||||
def test_init_with_missing_extension_dependency(self, extension_mock):
|
||||
extension_mock.side_effect = ImportError
|
||||
builder = PromptBuilder(template="This is a {{ variable }}")
|
||||
assert builder._env.extensions == {}
|
||||
res = builder.run(variable="test")
|
||||
assert res == {"prompt": "This is a test"}
|
||||
|
||||
def test_with_custom_dateformat(self) -> None:
|
||||
template = "Formatted date: {% now 'UTC', '%Y-%m-%d' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
result = builder.run()["prompt"]
|
||||
|
||||
now_formatted = f"Formatted date: {arrow.now('UTC').strftime('%Y-%m-%d')}"
|
||||
|
||||
assert now_formatted == result
|
||||
|
||||
def test_with_different_timezone(self) -> None:
|
||||
template = "Current time in New York is: {% now 'America/New_York' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
result = builder.run()["prompt"]
|
||||
|
||||
now_ny = f"Current time in New York is: {arrow.now('America/New_York').strftime('%Y-%m-%d %H:%M:%S')}"
|
||||
|
||||
assert now_ny == result
|
||||
|
||||
def test_date_with_addition_offset(self) -> None:
|
||||
template = "Time after 2 hours is: {% now 'UTC' + 'hours=2' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
result = builder.run()["prompt"]
|
||||
|
||||
now_plus_2 = f"Time after 2 hours is: {(arrow.now('UTC').shift(hours=+2)).strftime('%Y-%m-%d %H:%M:%S')}"
|
||||
|
||||
assert now_plus_2 == result
|
||||
|
||||
def test_date_with_subtraction_offset(self) -> None:
|
||||
template = "Time after 12 days is: {% now 'UTC' - 'days=12' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
result = builder.run()["prompt"]
|
||||
|
||||
now_plus_2 = f"Time after 12 days is: {(arrow.now('UTC').shift(days=-12)).strftime('%Y-%m-%d %H:%M:%S')}"
|
||||
|
||||
assert now_plus_2 == result
|
||||
|
||||
def test_invalid_timezone(self) -> None:
|
||||
template = "Current time is: {% now 'Invalid/Timezone' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
# Expect ValueError for invalid timezone
|
||||
with pytest.raises(ValueError, match="Invalid timezone"):
|
||||
builder.run()
|
||||
|
||||
def test_invalid_offset(self) -> None:
|
||||
template = "Time after invalid offset is: {% now 'UTC' + 'invalid_offset' %}"
|
||||
builder = PromptBuilder(template=template)
|
||||
|
||||
# Expect ValueError for invalid offset
|
||||
with pytest.raises(ValueError, match="Invalid offset or operator"):
|
||||
builder.run()
|
||||
Reference in New Issue
Block a user