deepset-ai--haystack/test/components/generators/chat/test_llm.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any
from unittest.mock import MagicMock

import pytest

from haystack import Document, Pipeline, component
from haystack.components.agents.agent import Agent
from haystack.components.generators.chat import LLM
from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.joiners.branch import BranchJoiner
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.routers.conditional_router import ConditionalRouter
from haystack.core.component.types import InputSocket, OutputSocket
from haystack.dataclasses import ChatMessage
from haystack.dataclasses.chat_message import ChatRole
from haystack.dataclasses.streaming_chunk import StreamingChunk
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.tools import Tool
from haystack.tools.toolset import Toolset


def sync_streaming_callback(chunk: StreamingChunk) -> None:
    pass


@component
class MockChatGeneratorWithTools:
    """A mock chat generator that accepts a tools parameter."""

    def to_dict(self) -> dict[str, Any]:
        return {"type": "test_llm.MockChatGeneratorWithTools", "data": {}}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithTools":
        return cls()

    @component.output_types(replies=list[ChatMessage])
    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
        return {"replies": [ChatMessage.from_assistant("Reply with tools support")]}

    @component.output_types(replies=list[ChatMessage])
    async def run_async(
        self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
    ) -> dict[str, Any]:
        return {"replies": [ChatMessage.from_assistant("Async reply with tools support")]}


@component
class MockChatGenerator:
    """A mock chat generator that does NOT accept a tools parameter."""

    def to_dict(self) -> dict[str, Any]:
        return {"type": "test_llm.MockChatGenerator", "data": {}}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MockChatGenerator":
        return cls()

    @component.output_types(replies=list[ChatMessage])
    def run(self, messages: list[ChatMessage], **kwargs) -> dict[str, Any]:
        return {"replies": [ChatMessage.from_assistant("Sync reply")]}

    @component.output_types(replies=list[ChatMessage])
    async def run_async(self, messages: list[ChatMessage], **kwargs) -> dict[str, Any]:
        return {"replies": [ChatMessage.from_assistant("Async reply")]}


class TestLLM:
    class TestInit:
        USER_PROMPT = '{% message role="user" %}{{ query }}{% endmessage %}'

        def test_is_subclass_of_agent(self):
            assert issubclass(LLM, Agent)

        def test_defaults(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            assert llm.chat_generator is not None
            assert llm.tools == []
            assert llm.system_prompt is None
            assert llm.user_prompt == self.USER_PROMPT
            assert llm.required_variables == "*"
            assert llm.streaming_callback is None

        def test_output_sockets(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            assert llm.__haystack_output__._sockets_dict == {
                "messages": OutputSocket(name="messages", type=list[ChatMessage], receivers=[]),
                "last_message": OutputSocket(name="last_message", type=ChatMessage, receivers=[]),
                "token_usage": OutputSocket(name="token_usage", type=dict[str, Any], receivers=[]),
            }

        def test_detects_no_tools_support(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            assert llm._chat_generator_supports_tools is False

        def test_detects_tools_support(self):
            llm = LLM(chat_generator=MockChatGeneratorWithTools(), user_prompt=self.USER_PROMPT)
            assert llm._chat_generator_supports_tools is True

        def test_messages_required_when_no_prompt_variables(self):
            llm = LLM(
                chat_generator=MockChatGenerator(), user_prompt='{% message role="user" %}Hello world{% endmessage %}'
            )
            messages_socket = llm.__haystack_input__._sockets_dict["messages"]
            assert isinstance(messages_socket, InputSocket)
            assert messages_socket.is_mandatory

        def test_messages_optional_when_prompt_has_variables(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            messages_socket = llm.__haystack_input__._sockets_dict["messages"]
            assert isinstance(messages_socket, InputSocket)
            assert not messages_socket.is_mandatory

        def test_messages_optional_when_plain_prompt_has_variables(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt="Question: {{ query }}")
            messages_socket = llm.__haystack_input__._sockets_dict["messages"]
            assert isinstance(messages_socket, InputSocket)
            assert not messages_socket.is_mandatory
            assert "query" in llm.__haystack_input__._sockets_dict

        def test_runtime_prompt_overrides_not_component_inputs(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            assert "system_prompt" not in llm.__haystack_input__._sockets_dict
            assert "user_prompt" not in llm.__haystack_input__._sockets_dict

        def test_raises_if_required_variables_empty(self):
            with pytest.raises(ValueError, match="required_variables must not be empty"):
                LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT, required_variables=[])

    class TestSerialization:
        def test_to_dict_excludes_agent_only_params(self, monkeypatch):
            monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
            user_prompt = '{% message role="user" %}{{ query }}{% endmessage %}'
            llm = LLM(chat_generator=OpenAIChatGenerator(), system_prompt="You are helpful.", user_prompt=user_prompt)

            serialized = llm.to_dict()

            assert serialized["type"] == "haystack.components.generators.chat.llm.LLM"
            assert "chat_generator" in serialized["init_parameters"]
            assert serialized["init_parameters"]["system_prompt"] == "You are helpful."

            agent_only_params = [
                "tools",
                "exit_conditions",
                "max_agent_steps",
                "raise_on_tool_invocation_failure",
                "tool_concurrency_limit",
                "tool_streaming_callback_passthrough",
                "confirmation_strategies",
                "state_schema",
            ]
            for param in agent_only_params:
                assert param not in serialized["init_parameters"], (
                    f"Agent-only param '{param}' should not be serialized"
                )

        def test_to_dict_includes_llm_params(self, monkeypatch):
            monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
            llm = LLM(
                chat_generator=OpenAIChatGenerator(),
                system_prompt="Be concise.",
                user_prompt='{% message role="user" %}{{ query }}{% endmessage %}',
                required_variables=["query"],
            )

            serialized = llm.to_dict()

            assert serialized["init_parameters"]["system_prompt"] == "Be concise."
            assert "{{ query }}" in serialized["init_parameters"]["user_prompt"]
            assert serialized["init_parameters"]["required_variables"] == ["query"]
            assert serialized["init_parameters"]["streaming_callback"] is None

        def test_from_dict(self, monkeypatch):
            monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
            data = {
                "type": "haystack.components.generators.chat.llm.LLM",
                "init_parameters": {
                    "chat_generator": {
                        "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
                        "init_parameters": {
                            "model": "gpt-4o-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,
                        },
                    },
                    "system_prompt": "You are helpful.",
                    "user_prompt": '{% message role="user" %}{{ query }}{% endmessage %}',
                    "required_variables": "*",
                    "streaming_callback": None,
                },
            }

            llm = LLM.from_dict(data)

            assert isinstance(llm, LLM)
            assert isinstance(llm.chat_generator, OpenAIChatGenerator)
            assert llm.system_prompt == "You are helpful."
            assert llm.tools == []

        def test_roundtrip(self, monkeypatch):
            monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
            user_prompt = '{% message role="user" %}{{ query }}{% endmessage %}'
            original = LLM(
                chat_generator=OpenAIChatGenerator(), system_prompt="You are a poet.", user_prompt=user_prompt
            )

            restored = LLM.from_dict(original.to_dict())

            assert isinstance(restored, LLM)
            assert isinstance(restored.chat_generator, OpenAIChatGenerator)
            assert restored.system_prompt == original.system_prompt
            assert restored.tools == []

    class TestRun:
        USER_PROMPT = '{% message role="user" %}{{ query }}{% endmessage %}'

        def test_run_accepts_messages_via_kwargs(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            prior_message = ChatMessage.from_user("Some prior context")
            result = llm.run(query="What is 2+2?", messages=[prior_message])
            assert result["last_message"].text == "Sync reply"
            assert prior_message in result["messages"]

        def test_run_without_messages(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            result = llm.run(query="What is 2+2?")
            assert result["last_message"].text == "Sync reply"
            user_messages = [m for m in result["messages"] if m.is_from(ChatRole.USER)]
            assert any("What is 2+2?" in m.text for m in user_messages)

        def test_run_with_plain_user_prompt(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt="Question: {{ query }}")
            result = llm.run(query="What is 2+2?")
            assert result["last_message"].text == "Sync reply"
            user_messages = [m for m in result["messages"] if m.is_from(ChatRole.USER)]
            assert any("Question: What is 2+2?" in m.text for m in user_messages)

        @pytest.mark.asyncio
        async def test_run_async_accepts_messages_via_kwargs(self):
            llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
            prior_message = ChatMessage.from_user("Some prior context")
            result = await llm.run_async(query="What is 2+2?", messages=[prior_message])
            assert result["last_message"].text == "Async reply"
            assert prior_message in result["messages"]

    class TestPipelineIntegration:
        @pytest.fixture()
        def document_store_with_docs(self):
            store = InMemoryDocumentStore()
            store.write_documents(
                [
                    Document(content="The Eiffel Tower is located in Paris."),
                    Document(content="The Brandenburg Gate is in Berlin."),
                    Document(content="The Colosseum is in Rome."),
                ]
            )
            return store

        def test_rag_pipeline(self, document_store_with_docs):
            user_prompt = (
                '{% message role="user" %}'
                "Use the following documents to answer the question.\n"
                "Documents:\n{% for doc in documents %}{{ doc.content }}\n{% endfor %}"
                "Question: {{ query }}"
                "{% endmessage %}"
            )
            llm = LLM(
                chat_generator=MockChatGenerator(),
                system_prompt="You are a knowledgeable assistant.",
                user_prompt=user_prompt,
                required_variables=["query", "documents"],
            )

            pipe = Pipeline()
            pipe.add_component("retriever", InMemoryBM25Retriever(document_store=document_store_with_docs))
            pipe.add_component("llm", llm)
            pipe.connect("retriever.documents", "llm.documents")

            query = "Where is the Colosseum?"
            result = pipe.run(data={"retriever": {"query": query}, "llm": {"query": query}})

            assert "llm" in result
            llm_output = result["llm"]
            assert "messages" in llm_output
            assert "last_message" in llm_output

            messages = llm_output["messages"]

            assert messages[0].is_from(ChatRole.SYSTEM)
            assert messages[0].text == "You are a knowledgeable assistant."

            user_messages = [m for m in messages if m.is_from(ChatRole.USER)]
            assert len(user_messages) == 1
            rendered = user_messages[0].text
            assert "Question: Where is the Colosseum?" in rendered
            assert "Documents:" in rendered
            assert "Colosseum" in rendered

            assert llm_output["last_message"].is_from(ChatRole.ASSISTANT)
            assert llm_output["last_message"].text == "Sync reply"


class TestLLMNotTriggeredByInjectedInput:
    """
    Regression guard for the optional-messages scheduling hazard described in
    https://github.com/deepset-ai/haystack/issues/11109.

    When `user_prompt` contains template variables, `messages` is optional on the LLM.
    An optional input with `sender=None` (i.e., injected directly via `pipeline.run`)
    would flip `has_user_input()` to True and incorrectly trigger the component even
    when its required inputs (e.g. `query`) never arrive.
    """

    def test_llm_not_triggered_by_injected_streaming_callback(self):

        @component
        class Planner:
            @component.output_types(messages=list[ChatMessage], last_role=str)
            def run(self) -> dict:
                return {"messages": [ChatMessage.from_user("hello")], "last_role": "assistant"}

        chat_generator = MockChatGenerator()
        llm = LLM(chat_generator=chat_generator)
        chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("x")]})

        router = ConditionalRouter(
            routes=[
                {
                    "condition": "{{ last_role == 'tool' }}",
                    "output": "{{ messages }}",
                    "output_name": "processing",
                    "output_type": list[ChatMessage],
                },
                {
                    "condition": "{{ True }}",
                    "output": "{{ messages }}",
                    "output_name": "planning",
                    "output_type": list[ChatMessage],
                },
            ],
            unsafe=True,
        )

        pipeline = Pipeline()
        pipeline.add_component("planner", Planner())
        pipeline.add_component("router", router)
        pipeline.add_component("branch_joiner", BranchJoiner(type_=list[ChatMessage]))
        pipeline.add_component("llm", llm)
        pipeline.connect("planner.messages", "router.messages")
        pipeline.connect("planner.last_role", "router.last_role")
        pipeline.connect("router.processing", "branch_joiner.value")
        pipeline.connect("branch_joiner.value", "llm.messages")

        result = pipeline.run(data={"llm": {"streaming_callback": sync_streaming_callback}})

        assert "llm" not in result
        chat_generator.run.assert_not_called()