haystack-experimental 0.14.2__py3-none-any.whl → 0.15.0__py3-none-any.whl

haystack_experimental/components/agents/agent.py

@@ -22,11 +22,11 @@ import haystack_experimental.core.pipeline.breakpoint as exp_breakpoint
 hs_breakpoint._create_agent_snapshot = exp_breakpoint._create_agent_snapshot
 hs_breakpoint._create_pipeline_snapshot_from_tool_invoker = exp_breakpoint._create_pipeline_snapshot_from_tool_invoker  # type: ignore[assignment]
 
-from haystack import logging
+from haystack import DeserializationError, logging
 from haystack.components.agents.agent import Agent as HaystackAgent
 from haystack.components.agents.agent import _ExecutionContext as Haystack_ExecutionContext
 from haystack.components.agents.agent import _schema_from_dict
-from haystack.components.agents.state import replace_values
+from haystack.components.agents.state import replace_values, State
 from haystack.components.generators.chat.types import ChatGenerator
 from haystack.core.errors import PipelineRuntimeError
 from haystack.core.pipeline import AsyncPipeline, Pipeline
@@ -36,19 +36,25 @@ from haystack.core.pipeline.breakpoint import (
 )
 from haystack.core.pipeline.utils import _deepcopy_with_exceptions
 from haystack.core.serialization import default_from_dict, import_class_by_name
-from haystack.dataclasses import ChatMessage
+from haystack.dataclasses import ChatMessage, ChatRole
 from haystack.dataclasses.breakpoints import AgentBreakpoint, ToolBreakpoint
-from haystack.dataclasses.streaming_chunk import StreamingCallbackT
+from haystack.dataclasses.streaming_chunk import StreamingCallbackT, select_streaming_callback
 from haystack.tools import ToolsType, deserialize_tools_or_toolset_inplace
 from haystack.utils.callable_serialization import deserialize_callable
 from haystack.utils.deserialization import deserialize_chatgenerator_inplace
 
+from haystack_experimental.chat_message_stores.types import ChatMessageStore
 from haystack_experimental.components.agents.human_in_the_loop import (
     ConfirmationStrategy,
     ToolExecutionDecision,
     HITLBreakpointException,
 )
-from haystack_experimental.components.agents.human_in_the_loop.strategies import _process_confirmation_strategies
+from haystack_experimental.components.agents.human_in_the_loop.strategies import (
+    _process_confirmation_strategies,
+    _process_confirmation_strategies_async,
+)
+from haystack_experimental.components.retrievers import ChatMessageRetriever
+from haystack_experimental.components.writers import ChatMessageWriter
 
 logger = logging.getLogger(__name__)
 
@@ -62,9 +68,15 @@ class _ExecutionContext(Haystack_ExecutionContext):
 
     :param tool_execution_decisions: Optional list of ToolExecutionDecision objects to use instead of prompting
         the user. This is useful when restarting from a snapshot where tool execution decisions were already made.
+    :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
+        to confirmation strategies. In web/server environments, this enables passing per-request
+        objects (e.g., WebSocket connections, async queues, or pub/sub clients) that strategies can use for
+        non-blocking user interaction. This is passed directly to strategies via the `confirmation_strategy_context`
+        parameter in their `run()` and `run_async()` methods.
     """
 
     tool_execution_decisions: Optional[list[ToolExecutionDecision]] = None
+    confirmation_strategy_context: Optional[dict[str, Any]] = None
 
 
 class Agent(HaystackAgent):
@@ -131,6 +143,7 @@ class Agent(HaystackAgent):
         raise_on_tool_invocation_failure: bool = False,
         confirmation_strategies: Optional[dict[str, ConfirmationStrategy]] = None,
         tool_invoker_kwargs: Optional[dict[str, Any]] = None,
+        chat_message_store: Optional[ChatMessageStore] = None,
     ) -> None:
         """
         Initialize the agent component.
@@ -164,6 +177,13 @@ class Agent(HaystackAgent):
             tool_invoker_kwargs=tool_invoker_kwargs,
         )
         self._confirmation_strategies = confirmation_strategies or {}
+        self._chat_message_store = chat_message_store
+        self._chat_message_retriever = (
+            ChatMessageRetriever(chat_message_store=chat_message_store) if chat_message_store else None
+        )
+        self._chat_message_writer = (
+            ChatMessageWriter(chat_message_store=chat_message_store) if chat_message_store else None
+        )
 
     def _initialize_fresh_execution(
         self,
@@ -172,7 +192,10 @@ class Agent(HaystackAgent):
         requires_async: bool,
         *,
         system_prompt: Optional[str] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         tools: Optional[Union[ToolsType, list[str]]] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
+        chat_message_store_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: dict[str, Any],
     ) -> _ExecutionContext:
         """
@@ -184,27 +207,56 @@ class Agent(HaystackAgent):
         :param system_prompt: System prompt for the agent. If provided, it overrides the default system prompt.
         :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
             When passing tool names, tools are selected from the Agent's originally configured tools.
+        :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
+            to confirmation strategies.
+        :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
         :param kwargs: Additional data to pass to the State used by the Agent.
         """
-        exe_context = super(Agent, self)._initialize_fresh_execution(
-            messages=messages,
-            streaming_callback=streaming_callback,
-            requires_async=requires_async,
-            system_prompt=system_prompt,
-            tools=tools,
-            **kwargs,
+        system_prompt = system_prompt or self.system_prompt
+        if system_prompt is not None:
+            messages = [ChatMessage.from_system(system_prompt)] + messages
+
+        # NOTE: difference with parent method to add chat message retrieval
+        if self._chat_message_retriever:
+            retriever_kwargs = _select_kwargs(self._chat_message_retriever, chat_message_store_kwargs or {})
+            if "chat_history_id" in retriever_kwargs:
+                messages = self._chat_message_retriever.run(
+                    current_messages=messages,
+                    **retriever_kwargs,
+                )["messages"]
+
+        if all(m.is_from(ChatRole.SYSTEM) for m in messages):
+            logger.warning("All messages provided to the Agent component are system messages. This is not recommended.")
+
+        state = State(schema=self.state_schema, data=kwargs)
+        state.set("messages", messages)
+
+        streaming_callback = select_streaming_callback(  # type: ignore[call-overload]
+            init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async
         )
-        # NOTE: 1st difference with parent method to add this to tool_invoker_inputs
+
+        selected_tools = self._select_tools(tools)
+        tool_invoker_inputs: dict[str, Any] = {"tools": selected_tools}
+        generator_inputs: dict[str, Any] = {"tools": selected_tools}
+        if streaming_callback is not None:
+            tool_invoker_inputs["streaming_callback"] = streaming_callback
+            generator_inputs["streaming_callback"] = streaming_callback
+        if generation_kwargs is not None:
+            generator_inputs["generation_kwargs"] = generation_kwargs
+
+        # NOTE: difference with parent method to add this to tool_invoker_inputs
         if self._tool_invoker:
-            exe_context.tool_invoker_inputs["enable_streaming_callback_passthrough"] = (
+            tool_invoker_inputs["enable_streaming_callback_passthrough"] = (
                 self._tool_invoker.enable_streaming_callback_passthrough
             )
-        # NOTE: 2nd difference is to use the extended _ExecutionContext
+
+        # NOTE: difference is to use the extended _ExecutionContext with confirmation_strategy_context
         return _ExecutionContext(
-            state=exe_context.state,
-            component_visits=exe_context.component_visits,
-            chat_generator_inputs=exe_context.chat_generator_inputs,
-            tool_invoker_inputs=exe_context.tool_invoker_inputs,
+            state=state,
+            component_visits=dict.fromkeys(["chat_generator", "tool_invoker"], 0),
+            chat_generator_inputs=generator_inputs,
+            tool_invoker_inputs=tool_invoker_inputs,
+            confirmation_strategy_context=confirmation_strategy_context,
         )
 
     def _initialize_from_snapshot(  # type: ignore[override]
@@ -213,7 +265,9 @@ class Agent(HaystackAgent):
         streaming_callback: Optional[StreamingCallbackT],
         requires_async: bool,
         *,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         tools: Optional[Union[ToolsType, list[str]]] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
     ) -> _ExecutionContext:
         """
         Initialize execution context from an AgentSnapshot.
@@ -221,18 +275,35 @@ class Agent(HaystackAgent):
         :param snapshot: An AgentSnapshot containing the state of a previously saved agent execution.
         :param streaming_callback: Optional callback for streaming responses.
         :param requires_async: Whether the agent run requires asynchronous execution.
+        :param generation_kwargs: Additional keyword arguments for chat generator. These parameters will
+            override the parameters passed during component initialization.
         :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
             When passing tool names, tools are selected from the Agent's originally configured tools.
+        :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
+            to confirmation strategies.
         """
-        exe_context = super(Agent, self)._initialize_from_snapshot(
-            snapshot=snapshot, streaming_callback=streaming_callback, requires_async=requires_async, tools=tools
-        )
+        # The PR https://github.com/deepset-ai/haystack/pull/9616 added the generation_kwargs parameter to
+        # _initialize_from_snapshot. This change has been released in Haystack 2.20.0.
+        # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
+        if inspect.signature(super(Agent, self)._initialize_from_snapshot).parameters.get("generation_kwargs"):
+            exe_context = super(Agent, self)._initialize_from_snapshot(  # type: ignore[call-arg]
+                snapshot=snapshot,
+                streaming_callback=streaming_callback,
+                requires_async=requires_async,
+                generation_kwargs=generation_kwargs,
+                tools=tools,
+            )
+        else:
+            exe_context = super(Agent, self)._initialize_from_snapshot(
+                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=requires_async, tools=tools
+            )
         # NOTE: 1st difference with parent method to add this to tool_invoker_inputs
         if self._tool_invoker:
             exe_context.tool_invoker_inputs["enable_streaming_callback_passthrough"] = (
                 self._tool_invoker.enable_streaming_callback_passthrough
             )
-        # NOTE: 2nd difference is to use the extended _ExecutionContext and add tool_execution_decisions
+        # NOTE: 2nd difference is to use the extended _ExecutionContext
+        # and add tool_execution_decisions + confirmation_strategy_context
         return _ExecutionContext(
             state=exe_context.state,
             component_visits=exe_context.component_visits,
@@ -241,17 +312,21 @@ class Agent(HaystackAgent):
             counter=exe_context.counter,
             skip_chat_generator=exe_context.skip_chat_generator,
             tool_execution_decisions=snapshot.tool_execution_decisions,
+            confirmation_strategy_context=confirmation_strategy_context,
         )
 
-    def run(  # noqa: PLR0915
+    def run(  # type: ignore[override]  # noqa: PLR0915
         self,
         messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         *,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         break_point: Optional[AgentBreakpoint] = None,
-        snapshot: Optional[AgentSnapshot] = None,  # type: ignore[override]
+        snapshot: Optional[AgentSnapshot] = None,
         system_prompt: Optional[str] = None,
         tools: Optional[Union[ToolsType, list[str]]] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
+        chat_message_store_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """
@@ -260,6 +335,8 @@ class Agent(HaystackAgent):
         :param messages: List of Haystack ChatMessage objects to process.
         :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
             The same callback can be configured to emit tool results when a tool is called.
+        :param generation_kwargs: Additional keyword arguments for LLM. These parameters will
+            override the parameters passed during component initialization.
         :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
             for "tool_invoker".
         :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
@@ -267,6 +344,12 @@ class Agent(HaystackAgent):
         :param system_prompt: System prompt for the agent. If provided, it overrides the default system prompt.
         :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
             When passing tool names, tools are selected from the Agent's originally configured tools.
+        :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
+            to confirmation strategies. Useful in web/server environments to provide per-request
+            objects (e.g., WebSocket connections, async queues, Redis pub/sub clients) that strategies
+            can use for non-blocking user interaction.
+        :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
+            For example, it can include the `chat_history_id` and `last_k` parameters for retrieving chat history.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -290,13 +373,18 @@ class Agent(HaystackAgent):
         # _runtime_checks. This change will be released in Haystack 2.20.0.
         # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
         if len(inspect.signature(self._runtime_checks).parameters) == 2:
-            self._runtime_checks(break_point, snapshot)
+            self._runtime_checks(break_point, snapshot)  # type: ignore[call-arg]  # pylint: disable=too-many-function-args
         else:
             self._runtime_checks(break_point)  # type: ignore[call-arg]  # pylint: disable=no-value-for-parameter
 
         if snapshot:
             exe_context = self._initialize_from_snapshot(
-                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=False, tools=tools
+                snapshot=snapshot,
+                streaming_callback=streaming_callback,
+                requires_async=False,
+                generation_kwargs=generation_kwargs,
+                tools=tools,
+                confirmation_strategy_context=confirmation_strategy_context,
             )
         else:
             exe_context = self._initialize_fresh_execution(
@@ -304,7 +392,10 @@ class Agent(HaystackAgent):
                 streaming_callback=streaming_callback,
                 requires_async=False,
                 system_prompt=system_prompt,
+                generation_kwargs=generation_kwargs,
                 tools=tools,
+                confirmation_strategy_context=confirmation_strategy_context,
+                chat_message_store_kwargs=chat_message_store_kwargs,
                 **kwargs,
             )
 
@@ -431,17 +522,27 @@ class Agent(HaystackAgent):
         result = {**exe_context.state.data}
         if msgs := result.get("messages"):
             result["last_message"] = msgs[-1]
+
+        # Write messages to ChatMessageStore if configured
+        if self._chat_message_writer:
+            writer_kwargs = _select_kwargs(self._chat_message_writer, chat_message_store_kwargs or {})
+            if "chat_history_id" in writer_kwargs:
+                self._chat_message_writer.run(messages=result["messages"], **writer_kwargs)
+
         return result
 
-    async def run_async(
+    async def run_async(  # type: ignore[override]
         self,
         messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         *,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         break_point: Optional[AgentBreakpoint] = None,
-        snapshot: Optional[AgentSnapshot] = None,  # type: ignore[override]
+        snapshot: Optional[AgentSnapshot] = None,
         system_prompt: Optional[str] = None,
         tools: Optional[Union[ToolsType, list[str]]] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
+        chat_message_store_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """
@@ -454,12 +555,20 @@ class Agent(HaystackAgent):
         :param messages: List of Haystack ChatMessage objects to process.
         :param streaming_callback: An asynchronous callback that will be invoked when a response is streamed from the
             LLM. The same callback can be configured to emit tool results when a tool is called.
+        :param generation_kwargs: Additional keyword arguments for LLM. These parameters will
+            override the parameters passed during component initialization.
         :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
             for "tool_invoker".
         :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
             the relevant information to restart the Agent execution from where it left off.
         :param system_prompt: System prompt for the agent. If provided, it overrides the default system prompt.
         :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
+        :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
+            to confirmation strategies. Useful in web/server environments to provide per-request
+            objects (e.g., WebSocket connections, async queues, Redis pub/sub clients) that strategies
+            can use for non-blocking user interaction.
+        :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
+            For example, it can include the `chat_history_id` and `last_k` parameters for retrieving chat history.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -483,13 +592,18 @@ class Agent(HaystackAgent):
         # _runtime_checks. This change will be released in Haystack 2.20.0.
         # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
         if len(inspect.signature(self._runtime_checks).parameters) == 2:
-            self._runtime_checks(break_point, snapshot)
+            self._runtime_checks(break_point, snapshot)  # type: ignore[call-arg]  # pylint: disable=too-many-function-args
         else:
             self._runtime_checks(break_point)  # type: ignore[call-arg]  # pylint: disable=no-value-for-parameter
 
         if snapshot:
             exe_context = self._initialize_from_snapshot(
-                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=True, tools=tools
+                snapshot=snapshot,
+                streaming_callback=streaming_callback,
+                requires_async=True,
+                generation_kwargs=generation_kwargs,
+                tools=tools,
+                confirmation_strategy_context=confirmation_strategy_context,
             )
         else:
             exe_context = self._initialize_fresh_execution(
@@ -497,7 +611,10 @@ class Agent(HaystackAgent):
                 streaming_callback=streaming_callback,
                 requires_async=True,
                 system_prompt=system_prompt,
+                generation_kwargs=generation_kwargs,
                 tools=tools,
+                confirmation_strategy_context=confirmation_strategy_context,
+                chat_message_store_kwargs=chat_message_store_kwargs,
                 **kwargs,
             )
 
@@ -535,8 +652,8 @@ class Agent(HaystackAgent):
 
                 # Apply confirmation strategies and update State and messages sent to ToolInvoker
                 try:
-                    # Run confirmation strategies to get updated tool call messages and modified chat history
-                    modified_tool_call_messages, new_chat_history = _process_confirmation_strategies(
+                    # Run confirmation strategies to get updated tool call messages and modified chat history (async)
+                    modified_tool_call_messages, new_chat_history = await _process_confirmation_strategies_async(
                         confirmation_strategies=self._confirmation_strategies,
                         messages_with_tool_calls=llm_messages,
                         execution_context=exe_context,
@@ -600,6 +717,13 @@ class Agent(HaystackAgent):
         result = {**exe_context.state.data}
         if msgs := result.get("messages"):
             result["last_message"] = msgs[-1]
+
+        # Write messages to ChatMessageStore if configured
+        if self._chat_message_writer:
+            writer_kwargs = _select_kwargs(self._chat_message_writer, chat_message_store_kwargs or {})
+            if "chat_history_id" in writer_kwargs:
+                self._chat_message_writer.run(messages=result["messages"], **writer_kwargs)
+
         return result
 
     def to_dict(self) -> dict[str, Any]:
@@ -614,6 +738,9 @@ class Agent(HaystackAgent):
             if self._confirmation_strategies
             else None
         )
+        data["init_parameters"]["chat_message_store"] = (
+            self._chat_message_store.to_dict() if self._chat_message_store is not None else None
+        )
         return data
 
     @classmethod
@@ -638,9 +765,31 @@ class Agent(HaystackAgent):
 
         if "confirmation_strategies" in init_params and init_params["confirmation_strategies"] is not None:
             for name, strategy_dict in init_params["confirmation_strategies"].items():
-                strategy_class = import_class_by_name(strategy_dict["type"])
+                try:
+                    strategy_class = import_class_by_name(strategy_dict["type"])
+                except ImportError as e:
+                    raise DeserializationError(f"Class '{strategy_dict['type']}' not correctly imported") from e
                 if not hasattr(strategy_class, "from_dict"):
-                    raise TypeError(f"{strategy_class} does not have from_dict method implemented.")
+                    raise DeserializationError(f"{strategy_class} does not have from_dict method implemented.")
                 init_params["confirmation_strategies"][name] = strategy_class.from_dict(strategy_dict)
 
+        if "chat_message_store" in init_params and init_params["chat_message_store"] is not None:
+            cms_data = init_params["chat_message_store"]
+            try:
+                cms_class = import_class_by_name(cms_data["type"])
+            except ImportError as e:
+                raise DeserializationError(f"Class '{cms_data['type']}' not correctly imported") from e
+            if not hasattr(cms_class, "from_dict"):
+                raise DeserializationError(f"{cms_class} does not have from_dict method implemented.")
+            init_params["chat_message_store"] = cms_class.from_dict(cms_data)
+
         return default_from_dict(cls, data)
+
+
+def _select_kwargs(obj: Any, source: dict) -> dict[str, Any]:
+    """
+    Select only those key-value pairs from source dict that are valid parameters for obj.run() method.
+    """
+    sig = inspect.signature(obj.run)
+    allowed = set(sig.parameters.keys())
+    return {k: v for k, v in source.items() if k in allowed}