# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import logging import os from unittest.mock import AsyncMock, Mock import pytest from haystack.components.generators.chat.openai import OpenAIChatGenerator from haystack.components.query.query_expander import DEFAULT_PROMPT_TEMPLATE, QueryExpander from haystack.dataclasses.chat_message import ChatMessage @pytest.fixture def mock_chat_generator(): return Mock(spec=OpenAIChatGenerator) class TestQueryExpander: def test_init_default_generator(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() assert expander.n_expansions == 4 assert expander.include_original_query is True assert isinstance(expander.chat_generator, OpenAIChatGenerator) assert expander.chat_generator.model == "gpt-4.1-mini" assert expander._prompt_builder is not None def test_init_custom_generator(self, mock_chat_generator): expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3) assert expander.n_expansions == 3 assert expander.chat_generator is mock_chat_generator def test_run_warms_up_chat_generator(self, mock_chat_generator): expander = QueryExpander(chat_generator=mock_chat_generator) mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("1. test query")]} expander.run("test query") mock_chat_generator.warm_up.assert_called() def test_init_negative_expansions_raises_error(self): with pytest.raises(ValueError, match="n_expansions must be positive"): QueryExpander(n_expansions=-1) def test_init_zero_expansions_raises_error(self): with pytest.raises(ValueError, match="n_expansions must be positive"): QueryExpander(n_expansions=0) def test_init_custom_prompt_template(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") custom_template = "Custom template: {{ query }} with {{ n_expansions }} expansions" expander = QueryExpander(prompt_template=custom_template) assert expander.prompt_template == custom_template def test_run_negative_expansions_raises_error(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() with pytest.raises(ValueError, match="n_expansions must be positive"): expander.run("test query", n_expansions=-1) def test_run_zero_expansions_raises_error(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander(n_expansions=4) with pytest.raises(ValueError, match="n_expansions must be positive"): expander.run("test query", n_expansions=0) def test_run_with_runtime_n_expansions_override(self, mock_chat_generator): mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["alt1", "alt2"]}')] } expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=4, include_original_query=False) result = expander.run("test query", n_expansions=2) # should request 2 expansions call_args = mock_chat_generator.run.call_args[1]["messages"][0].text assert "2" in call_args assert len(result["queries"]) == 2 assert result["queries"] == ["alt1", "alt2"] def test_run_successful_expansion(self, mock_chat_generator): mock_chat_generator.run.return_value = { "replies": [ ChatMessage.from_assistant( '{"queries": ["alternative query 1", "alternative query 2", "alternative query 3"]}' ) ] } expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3) result = expander.run("original query") assert result["queries"] == [ "alternative query 1", "alternative query 2", "alternative query 3", "original query", ] mock_chat_generator.run.assert_called_once() def test_run_without_including_original(self, mock_chat_generator): mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["alt1", "alt2"]}')] } expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=False) result = expander.run("original") assert result["queries"] == ["alt1", "alt2"] def test_run_empty_query(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() result = expander.run("") assert result["queries"] == [""] def test_run_empty_query_no_original(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander(include_original_query=False) result = expander.run(" ") assert result["queries"] == [] def test_run_whitespace_only_query(self, monkeypatch, caplog): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() with caplog.at_level(logging.WARNING): result = expander.run("\t\n \r") assert result["queries"] == ["\t\n \r"] assert "Empty query provided" in caplog.text def test_run_generator_no_replies(self, mock_chat_generator): mock_chat_generator.run.return_value = {"replies": []} expander = QueryExpander(chat_generator=mock_chat_generator) result = expander.run("test query") assert result["queries"] == ["test query"] def test_run_generator_exception(self, mock_chat_generator): mock_chat_generator.run.side_effect = Exception("Generator error") expander = QueryExpander(chat_generator=mock_chat_generator) result = expander.run("test query") assert result["queries"] == ["test query"] def test_run_invalid_json_response(self, mock_chat_generator): mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("invalid json response")]} expander = QueryExpander(chat_generator=mock_chat_generator) result = expander.run("test query") assert result["queries"] == ["test query"] def test_parse_expanded_queries_valid_json(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() queries = expander._parse_expanded_queries('{"queries": ["query1", "query2", "query3"]}') assert queries == ["query1", "query2", "query3"] def test_parse_expanded_queries_invalid_json(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() queries = expander._parse_expanded_queries("not json") assert queries == [] def test_parse_expanded_queries_empty_string(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() queries = expander._parse_expanded_queries("") assert queries == [] def test_parse_expanded_queries_non_list_json(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() queries = expander._parse_expanded_queries('{"not": "a list"}') assert queries == [] @pytest.mark.parametrize("queries_value", ['"single query"', '{"query": "value"}', "null"]) def test_parse_expanded_queries_rejects_non_list_queries_value(self, monkeypatch, caplog, queries_value): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() with caplog.at_level(logging.WARNING): queries = expander._parse_expanded_queries(f'{{"queries": {queries_value}}}') assert queries == [] assert "Expected 'queries' to be a list" in caplog.text def test_parse_expanded_queries_mixed_types(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() queries = expander._parse_expanded_queries('{"queries": ["valid query", 123, "", "another valid"]}') assert queries == ["valid query", "another valid"] def test_run_query_deduplication(self, mock_chat_generator): mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["original query", "alt1", "alt2"]}')] } expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=True) result = expander.run("original query") assert result["queries"] == ["original query", "alt1", "alt2"] assert len(result["queries"]) == 3 def test_run_truncates_excess_queries(self, mock_chat_generator, caplog): mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["q1", "q2", "q3", "q4", "q5"]}')] } expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3, include_original_query=False) with caplog.at_level(logging.WARNING): result = expander.run("test query") assert len(result["queries"]) == 3 assert result["queries"] == ["q1", "q2", "q3"] assert "Generated 5 queries but only 3 were requested" in caplog.text assert "Truncating" in caplog.text def test_run_with_custom_template(self, mock_chat_generator): custom_template = """ Create {{ n_expansions }} alternative search queries for: {{ query }} Return as JSON: {"queries": ["query1", "query2"]} """ mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["custom alt 1", "custom alt 2"]}')] } expander = QueryExpander( chat_generator=mock_chat_generator, prompt_template=custom_template, n_expansions=2, include_original_query=False, ) result = expander.run("test query") assert result["queries"] == ["custom alt 1", "custom alt 2"] mock_chat_generator.run.assert_called_once() call_args = mock_chat_generator.run.call_args[1]["messages"][0].text assert "Create 2 alternative search queries for: test query" in call_args assert "Return as JSON" in call_args def test_component_output_types(self, mock_chat_generator, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") expander = QueryExpander() mock_chat_generator.run.return_value = { "replies": [ChatMessage.from_assistant('{"queries": ["test1", "test2"]}')] } expander.chat_generator = mock_chat_generator result = expander.run("test") assert "queries" in result assert isinstance(result["queries"], list) assert all(isinstance(q, str) for q in result["queries"]) @pytest.mark.parametrize("variable", ["query", "n_expansions"]) def test_prompt_template_missing_variable(self, caplog, variable, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") if variable == "query": template_missing_variable = "Generate {{ n_expansions }} expansions" else: template_missing_variable = "Generate expansions for {{ query }}" with caplog.at_level(logging.WARNING): QueryExpander(prompt_template=template_missing_variable) assert f"The prompt template does not contain the '{variable}' variable" in caplog.text assert "This may cause issues during execution" in caplog.text def test_to_dict(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") generator = OpenAIChatGenerator(model="gpt-4.1-mini") expander = QueryExpander(chat_generator=generator, n_expansions=2, include_original_query=False) serialized_query_expander = expander.to_dict() assert serialized_query_expander == { "type": "haystack.components.query.query_expander.QueryExpander", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4.1-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "prompt_template": DEFAULT_PROMPT_TEMPLATE, "n_expansions": 2, "include_original_query": False, }, } def test_from_dict(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345") data = { "type": "haystack.components.query.query_expander.QueryExpander", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4.1-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "prompt_template": DEFAULT_PROMPT_TEMPLATE, "n_expansions": 2, "include_original_query": False, }, } expander = QueryExpander.from_dict(data) assert expander.n_expansions == 2 assert expander.include_original_query is False assert expander.prompt_template == DEFAULT_PROMPT_TEMPLATE assert isinstance(expander.chat_generator, OpenAIChatGenerator) assert expander.chat_generator.model == "gpt-4.1-mini" class FakeSyncOnlyChatGenerator: """A chat generator exposing only a synchronous `run` (no `run_async`) for the fallback path.""" def __init__(self): self.run = Mock() class TestQueryExpanderAsync: @pytest.mark.asyncio async def test_run_async(self): mock_chat_generator = Mock(spec=OpenAIChatGenerator) mock_chat_generator.run_async = AsyncMock( return_value={ "replies": [ ChatMessage.from_assistant( '{"queries": ["alternative query 1", "alternative query 2", "alternative query 3"]}' ) ] } ) expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3) result = await expander.run_async("original query") assert result["queries"] == [ "alternative query 1", "alternative query 2", "alternative query 3", "original query", ] mock_chat_generator.run_async.assert_awaited_once() mock_chat_generator.run.assert_not_called() @pytest.mark.asyncio async def test_run_async_without_including_original(self): mock_chat_generator = Mock(spec=OpenAIChatGenerator) mock_chat_generator.run_async = AsyncMock( return_value={"replies": [ChatMessage.from_assistant('{"queries": ["alt1", "alt2"]}')]} ) expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=False) result = await expander.run_async("original") assert result["queries"] == ["alt1", "alt2"] mock_chat_generator.run_async.assert_awaited_once() @pytest.mark.asyncio async def test_run_async_fallback_to_sync_run(self): fake_chat_generator = FakeSyncOnlyChatGenerator() fake_chat_generator.run.return_value = { "replies": [ ChatMessage.from_assistant( '{"queries": ["alternative query 1", "alternative query 2", "alternative query 3"]}' ) ] } assert not hasattr(fake_chat_generator, "run_async") expander = QueryExpander(chat_generator=fake_chat_generator, n_expansions=3) result = await expander.run_async("original query") assert result["queries"] == [ "alternative query 1", "alternative query 2", "alternative query 3", "original query", ] fake_chat_generator.run.assert_called_once() @pytest.mark.integration class TestQueryExpanderIntegration: @pytest.fixture def chat_generator(self): return OpenAIChatGenerator( model="gpt-4.1-nano", generation_kwargs={ "temperature": 0.7, "response_format": { "type": "json_schema", "json_schema": { "name": "query_expansion", "schema": { "type": "object", "properties": { "queries": {"type": "array", "items": {"type": "string"}, "minItems": 2, "maxItems": 2} }, "required": ["queries"], "additionalProperties": False, }, }, }, }, ) @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_query_expansion(self, chat_generator): expander = QueryExpander(n_expansions=2, chat_generator=chat_generator) result = expander.run("renewable energy sources") assert len(result["queries"]) == 3 assert all(len(q.strip()) > 0 for q in result["queries"]) assert "renewable energy sources" in result["queries"] @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_query_expansion_async(self, chat_generator): expander = QueryExpander(n_expansions=2, chat_generator=chat_generator) result = await expander.run_async("renewable energy sources") assert len(result["queries"]) == 3 assert all(len(q.strip()) > 0 for q in result["queries"]) assert "renewable energy sources" in result["queries"] @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_different_domains(self, chat_generator): test_queries = ["machine learning algorithms", "climate change effects", "quantum computing applications"] expander = QueryExpander(n_expansions=2, include_original_query=False, chat_generator=chat_generator) for query in test_queries: result = expander.run(query) # Should return exactly 2 expansions (no original) assert len(result["queries"]) == 2 # Should be different from original assert query not in result["queries"] class TestComponentLifecycle: def test_warm_up_delegates_to_chat_generator(self, mock_chat_generator): expander = QueryExpander(chat_generator=mock_chat_generator) expander.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() expander = QueryExpander(chat_generator=mock_chat_generator) await expander.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"]) expander = QueryExpander(chat_generator=chat_generator) await expander.warm_up_async() chat_generator.warm_up.assert_called_once() def test_close_delegates_to_chat_generator(self, mock_chat_generator): expander = QueryExpander(chat_generator=mock_chat_generator) expander.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() expander = QueryExpander(chat_generator=mock_chat_generator) await expander.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"]) expander = QueryExpander(chat_generator=chat_generator) await expander.close_async() chat_generator.close.assert_called_once() def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): chat_generator = Mock(spec=["run"]) expander = QueryExpander(chat_generator=chat_generator) expander.warm_up() expander.close()