mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-07 00:51:22 +00:00
85258f0654
* Fix types in test_run.py * Get test_run.py to pass fmt-check * Add test_run to mypy checks * Update test folder to pass ruff linting * Fix merge * Fix HF tests * Fix hf test * Try to fix tests * Another attempt * minor fix * fix SentenceTransformersDiversityRanker * skip integrations tests due to model unavailable on HF inference --------- Co-authored-by: anakin87 <stefanofiorucci@gmail.com>
340 lines
14 KiB
Python
340 lines
14 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
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 is None
|
|
|
|
# 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 == Optional[str]
|
|
assert inputs["template_variables"].type == Optional[Dict[str, Any]]
|
|
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 == Optional[str]
|
|
assert inputs["template_variables"].type == Optional[Dict[str, Any]]
|
|
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 is None
|
|
|
|
# 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 == Optional[str]
|
|
assert inputs["template_variables"].type == Optional[Dict[str, Any]]
|
|
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
|
|
|
|
@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_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": None},
|
|
}
|
|
|
|
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_input(self):
|
|
builder = PromptBuilder(template="This is a {{ variable }}")
|
|
res = builder.run()
|
|
assert res == {"prompt": "This is a "}
|
|
|
|
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)
|
|
|
|
template_variables = {"name": "John"}
|
|
expected_result = {"prompt": "Hello, John! How are you?"}
|
|
|
|
assert builder.run(template_variables=template_variables, var1="How are you?") == expected_result
|
|
|
|
def test_run_overwriting_default_template(self):
|
|
default_template = "Hello, {{ name }}!"
|
|
|
|
builder = PromptBuilder(template=default_template)
|
|
|
|
template = "Hello, {{ var1 }}{{ name }}!"
|
|
expected_result = {"prompt": "Hello, John!"}
|
|
|
|
assert builder.run(template, name="John") == expected_result
|
|
|
|
def test_run_overwriting_default_template_with_template_variables(self):
|
|
default_template = "Hello, {{ name }}!"
|
|
|
|
builder = PromptBuilder(template=default_template)
|
|
|
|
template = "Hello, {{ var1 }} {{ name }}!"
|
|
template_variables = {"var1": "Big"}
|
|
expected_result = {"prompt": "Hello, Big John!"}
|
|
|
|
assert builder.run(template, template_variables, name="John") == expected_result
|
|
|
|
def test_run_overwriting_default_template_with_variables(self):
|
|
variables = ["var1", "var2", "name"]
|
|
default_template = "Hello, {{ name }}!"
|
|
|
|
builder = PromptBuilder(template=default_template, variables=variables)
|
|
|
|
template = "Hello, {{ var1 }} {{ name }}!"
|
|
expected_result = {"prompt": "Hello, Big John!"}
|
|
|
|
assert builder.run(template, name="John", var1="Big") == expected_result
|
|
|
|
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_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()
|
|
|
|
def test_warning_no_required_variables(self, caplog):
|
|
with caplog.at_level(logging.WARNING):
|
|
_ = PromptBuilder(template="This is a {{ variable }}")
|
|
assert "but `required_variables` is not set." in caplog.text
|