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

haystack_experimental/components/agents/human_in_the_loop/strategies.py

@@ -47,7 +47,13 @@ class BlockingConfirmationStrategy:
         self.confirmation_ui = confirmation_ui
 
     def run(
-        self, tool_name: str, tool_description: str, tool_params: dict[str, Any], tool_call_id: Optional[str] = None
+        self,
+        *,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
     ) -> ToolExecutionDecision:
         """
         Run the human-in-the-loop strategy for a given tool and its parameters.
@@ -61,6 +67,10 @@ class BlockingConfirmationStrategy:
         :param tool_call_id:
             Optional unique identifier for the tool call. This can be used to track and correlate the decision with a
             specific tool invocation.
+        :param confirmation_strategy_context:
+            Optional dictionary for passing request-scoped resources. 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.
 
         :returns:
             A ToolExecutionDecision indicating whether to execute the tool with the given parameters, or a
@@ -109,6 +119,40 @@ class BlockingConfirmationStrategy:
                 tool_name=tool_name, execute=True, tool_call_id=tool_call_id, final_tool_params=tool_params
             )
 
+    async def run_async(
+        self,
+        *,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
+    ) -> ToolExecutionDecision:
+        """
+        Async version of run. Calls the sync run() method by default.
+
+        :param tool_name:
+            The name of the tool to be executed.
+        :param tool_description:
+            The description of the tool.
+        :param tool_params:
+            The parameters to be passed to the tool.
+        :param tool_call_id:
+            Optional unique identifier for the tool call.
+        :param confirmation_strategy_context:
+            Optional dictionary for passing request-scoped resources.
+
+        :returns:
+            A ToolExecutionDecision indicating whether to execute the tool with the given parameters.
+        """
+        return self.run(
+            tool_name=tool_name,
+            tool_description=tool_description,
+            tool_params=tool_params,
+            tool_call_id=tool_call_id,
+            confirmation_strategy_context=confirmation_strategy_context,
+        )
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the BlockingConfirmationStrategy to a dictionary.
@@ -161,7 +205,13 @@ class BreakpointConfirmationStrategy:
         self.snapshot_file_path = snapshot_file_path
 
     def run(
-        self, tool_name: str, tool_description: str, tool_params: dict[str, Any], tool_call_id: Optional[str] = None
+        self,
+        *,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
     ) -> ToolExecutionDecision:
         """
         Run the breakpoint confirmation strategy for a given tool and its parameters.
@@ -175,6 +225,9 @@ class BreakpointConfirmationStrategy:
         :param tool_call_id:
             Optional unique identifier for the tool call. This can be used to track and correlate the decision with a
             specific tool invocation.
+        :param confirmation_strategy_context:
+            Optional dictionary for passing request-scoped resources. Not used by this strategy but included for
+            interface compatibility.
 
         :raises HITLBreakpointException:
             Always raises an `HITLBreakpointException` exception to signal that user confirmation is required.
@@ -189,6 +242,43 @@ class BreakpointConfirmationStrategy:
             snapshot_file_path=self.snapshot_file_path,
         )
 
+    async def run_async(
+        self,
+        *,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        confirmation_strategy_context: Optional[dict[str, Any]] = None,
+    ) -> ToolExecutionDecision:
+        """
+        Async version of run. Calls the sync run() method.
+
+        :param tool_name:
+            The name of the tool to be executed.
+        :param tool_description:
+            The description of the tool.
+        :param tool_params:
+            The parameters to be passed to the tool.
+        :param tool_call_id:
+            Optional unique identifier for the tool call.
+        :param confirmation_strategy_context:
+            Optional dictionary for passing request-scoped resources.
+
+        :raises HITLBreakpointException:
+            Always raises an `HITLBreakpointException` exception to signal that user confirmation is required.
+
+        :returns:
+            This method does not return; it always raises an exception.
+        """
+        return self.run(
+            tool_name=tool_name,
+            tool_description=tool_description,
+            tool_params=tool_params,
+            tool_call_id=tool_call_id,
+            confirmation_strategy_context=confirmation_strategy_context,
+        )
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the BreakpointConfirmationStrategy to a dictionary.
@@ -285,6 +375,46 @@ def _process_confirmation_strategies(
     return modified_tool_call_messages, new_chat_history
 
 
+async def _process_confirmation_strategies_async(
+    *,
+    confirmation_strategies: dict[str, ConfirmationStrategy],
+    messages_with_tool_calls: list[ChatMessage],
+    execution_context: "_ExecutionContext",
+) -> tuple[list[ChatMessage], list[ChatMessage]]:
+    """
+    Async version of _process_confirmation_strategies.
+
+    Run the confirmation strategies and return modified tool call messages and updated chat history.
+
+    :param confirmation_strategies: Mapping of tool names to their corresponding confirmation strategies
+    :param messages_with_tool_calls: Chat messages containing tool calls
+    :param execution_context: The current execution context of the agent
+    :returns:
+        Tuple of modified messages with confirmed tool calls and updated chat history
+    """
+    # Run confirmation strategies and get tool execution decisions (async version)
+    teds = await _run_confirmation_strategies_async(
+        confirmation_strategies=confirmation_strategies,
+        messages_with_tool_calls=messages_with_tool_calls,
+        execution_context=execution_context,
+    )
+
+    # Apply tool execution decisions to messages_with_tool_calls
+    rejection_messages, modified_tool_call_messages = _apply_tool_execution_decisions(
+        tool_call_messages=messages_with_tool_calls,
+        tool_execution_decisions=teds,
+    )
+
+    # Update the chat history with rejection messages and new tool call messages
+    new_chat_history = _update_chat_history(
+        chat_history=execution_context.state.get("messages"),
+        rejection_messages=rejection_messages,
+        tool_call_and_explanation_messages=modified_tool_call_messages,
+    )
+
+    return modified_tool_call_messages, new_chat_history
+
+
 def _run_confirmation_strategies(
     confirmation_strategies: dict[str, ConfirmationStrategy],
     messages_with_tool_calls: list[ChatMessage],
@@ -344,13 +474,100 @@ def _run_confirmation_strategies(
             # If not, run the confirmation strategy
             if not ted:
                 ted = confirmation_strategies[tool_name].run(
-                    tool_name=tool_name, tool_description=tool_to_invoke.description, tool_params=final_args
+                    tool_name=tool_name,
+                    tool_description=tool_to_invoke.description,
+                    tool_params=final_args,
+                    tool_call_id=tool_call.id,
+                    confirmation_strategy_context=execution_context.confirmation_strategy_context,
                 )
             teds.append(ted)
 
     return teds
 
 
+async def _run_confirmation_strategies_async(
+    confirmation_strategies: dict[str, ConfirmationStrategy],
+    messages_with_tool_calls: list[ChatMessage],
+    execution_context: "_ExecutionContext",
+) -> list[ToolExecutionDecision]:
+    """
+    Async version of _run_confirmation_strategies.
+
+    Run confirmation strategies for tool calls in the provided chat messages.
+
+    :param confirmation_strategies: Mapping of tool names to their corresponding confirmation strategies
+    :param messages_with_tool_calls: Messages containing tool calls to process
+    :param execution_context: The current execution context containing state and inputs
+    :returns:
+        A list of ToolExecutionDecision objects representing the decisions made for each tool call.
+    """
+    state = execution_context.state
+    tools_with_names = {tool.name: tool for tool in execution_context.tool_invoker_inputs["tools"]}
+    existing_teds = execution_context.tool_execution_decisions if execution_context.tool_execution_decisions else []
+    existing_teds_by_name = {ted.tool_name: ted for ted in existing_teds if ted.tool_name}
+    existing_teds_by_id = {ted.tool_call_id: ted for ted in existing_teds if ted.tool_call_id}
+
+    teds = []
+    for message in messages_with_tool_calls:
+        if not message.tool_calls:
+            continue
+
+        for tool_call in message.tool_calls:
+            tool_name = tool_call.tool_name
+            tool_to_invoke = tools_with_names[tool_name]
+
+            # Prepare final tool args
+            final_args = _prepare_tool_args(
+                tool=tool_to_invoke,
+                tool_call_arguments=tool_call.arguments,
+                state=state,
+                streaming_callback=execution_context.tool_invoker_inputs.get("streaming_callback"),
+                enable_streaming_passthrough=execution_context.tool_invoker_inputs.get(
+                    "enable_streaming_passthrough", False
+                ),
+            )
+
+            # Get tool execution decisions from confirmation strategies
+            # If no confirmation strategy is defined for this tool, proceed with execution
+            if tool_name not in confirmation_strategies:
+                teds.append(
+                    ToolExecutionDecision(
+                        tool_call_id=tool_call.id,
+                        tool_name=tool_name,
+                        execute=True,
+                        final_tool_params=final_args,
+                    )
+                )
+                continue
+
+            # Check if there's already a decision for this tool call in the execution context
+            ted = existing_teds_by_id.get(tool_call.id or "") or existing_teds_by_name.get(tool_name)
+
+            # If not, run the confirmation strategy (async version)
+            if not ted:
+                strategy = confirmation_strategies[tool_name]
+                # Use run_async if available, otherwise fall back to sync run
+                if hasattr(strategy, "run_async"):
+                    ted = await strategy.run_async(
+                        tool_name=tool_name,
+                        tool_description=tool_to_invoke.description,
+                        tool_params=final_args,
+                        tool_call_id=tool_call.id,
+                        confirmation_strategy_context=execution_context.confirmation_strategy_context,
+                    )
+                else:
+                    ted = strategy.run(
+                        tool_name=tool_name,
+                        tool_description=tool_to_invoke.description,
+                        tool_params=final_args,
+                        tool_call_id=tool_call.id,
+                        confirmation_strategy_context=execution_context.confirmation_strategy_context,
+                    )
+            teds.append(ted)
+
+    return teds
+
+
 def _apply_tool_execution_decisions(
     tool_call_messages: list[ChatMessage], tool_execution_decisions: list[ToolExecutionDecision]
 ) -> tuple[list[ChatMessage], list[ChatMessage]]:

haystack_experimental/components/agents/human_in_the_loop/types.py

@@ -63,7 +63,12 @@ class ConfirmationPolicy(Protocol):
 
 class ConfirmationStrategy(Protocol):
     def run(
-        self, tool_name: str, tool_description: str, tool_params: dict[str, Any], tool_call_id: Optional[str] = None
+        self,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        **kwargs: Optional[dict[str, Any]],
     ) -> ToolExecutionDecision:
         """
         Run the confirmation strategy for a given tool and its parameters.
@@ -73,6 +78,36 @@ class ConfirmationStrategy(Protocol):
         :param tool_params: The parameters to be passed to the tool.
         :param tool_call_id: Optional unique identifier for the tool call. This can be used to track and correlate
             the decision with a specific tool invocation.
+        :param kwargs: Additional keyword arguments. Implementations may accept `confirmation_strategy_context`
+            for passing request-scoped resources (e.g., WebSocket connections, async queues) in web/server
+            environments.
+
+        :returns:
+            The result of the confirmation strategy (e.g., tool output, rejection message, etc.).
+        """
+        ...
+
+    async def run_async(
+        self,
+        tool_name: str,
+        tool_description: str,
+        tool_params: dict[str, Any],
+        tool_call_id: Optional[str] = None,
+        **kwargs: Optional[dict[str, Any]],
+    ) -> ToolExecutionDecision:
+        """
+        Async version of run. Run the confirmation strategy for a given tool and its parameters.
+
+        Default implementation calls the sync run() method. Override for true async behavior.
+
+        :param tool_name: The name of the tool to be executed.
+        :param tool_description: The description of the tool.
+        :param tool_params: The parameters to be passed to the tool.
+        :param tool_call_id: Optional unique identifier for the tool call. This can be used to track and correlate
+            the decision with a specific tool invocation.
+        :param kwargs: Additional keyword arguments. Implementations may accept `confirmation_strategy_context`
+            for passing request-scoped resources (e.g., WebSocket connections, async queues) in web/server
+            environments.
 
         :returns:
             The result of the confirmation strategy (e.g., tool output, rejection message, etc.).

haystack_experimental/components/embedders/types/protocol.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, List, Protocol
+from typing import Any, Protocol
 
 from haystack import Document
 
@@ -15,7 +15,7 @@ class DocumentEmbedder(Protocol):
     Protocol for Document Embedders.
     """
 
-    def run(self, documents: List[Document]) -> Dict[str, Any]:
+    def run(self, documents: list[Document]) -> dict[str, Any]:
         """
         Generate embeddings for the input documents.
 

haystack_experimental/components/preprocessors/__init__.py

@@ -9,10 +9,12 @@ from lazy_imports import LazyImporter
 
 _import_structure = {
     "embedding_based_document_splitter": ["EmbeddingBasedDocumentSplitter"],
+    "md_header_level_inferrer": ["MarkdownHeaderLevelInferrer"],
 }
 
 if TYPE_CHECKING:
     from .embedding_based_document_splitter import EmbeddingBasedDocumentSplitter
+    from .md_header_level_inferrer import MarkdownHeaderLevelInferrer
 
 else:
     sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)

haystack_experimental/components/preprocessors/embedding_based_document_splitter.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 import numpy as np
 from haystack import Document, component, logging
@@ -136,8 +136,8 @@ class EmbeddingBasedDocumentSplitter:
             self.document_embedder.warm_up()
         self._is_warmed_up = True
 
-    @component.output_types(documents=List[Document])
-    def run(self, documents: List[Document]) -> Dict[str, List[Document]]:
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
         """
         Split documents based on embedding similarity.
 
@@ -162,7 +162,7 @@ class EmbeddingBasedDocumentSplitter:
         if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
             raise TypeError("EmbeddingBasedDocumentSplitter expects a List of Documents as input.")
 
-        split_docs: List[Document] = []
+        split_docs: list[Document] = []
         for doc in documents:
             if doc.content is None:
                 raise ValueError(
@@ -178,7 +178,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return {"documents": split_docs}
 
-    def _split_document(self, doc: Document) -> List[Document]:
+    def _split_document(self, doc: Document) -> list[Document]:
         """
         Split a single document based on embedding similarity.
         """
@@ -194,7 +194,7 @@ class EmbeddingBasedDocumentSplitter:
         # Create Document objects from the final splits
         return EmbeddingBasedDocumentSplitter._create_documents_from_splits(splits=final_splits, original_doc=doc)
 
-    def _split_text(self, text: str) -> List[str]:
+    def _split_text(self, text: str) -> list[str]:
         """
         Split a text into smaller chunks based on embedding similarity.
         """
@@ -221,7 +221,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return sub_splits
 
-    def _group_sentences(self, sentences: List[str]) -> List[str]:
+    def _group_sentences(self, sentences: list[str]) -> list[str]:
         """
         Group sentences into groups of sentences_per_group.
         """
@@ -235,7 +235,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return groups
 
-    def _calculate_embeddings(self, sentence_groups: List[str]) -> List[List[float]]:
+    def _calculate_embeddings(self, sentence_groups: list[str]) -> list[list[float]]:
         """
         Calculate embeddings for each sentence group using the DocumentEmbedder.
         """
@@ -246,7 +246,7 @@ class EmbeddingBasedDocumentSplitter:
         embeddings = [doc.embedding for doc in embedded_docs]
         return embeddings
 
-    def _find_split_points(self, embeddings: List[List[float]]) -> List[int]:
+    def _find_split_points(self, embeddings: list[list[float]]) -> list[int]:
         """
         Find split points based on cosine distances between sequential embeddings.
         """
@@ -273,7 +273,7 @@ class EmbeddingBasedDocumentSplitter:
         return split_points
 
     @staticmethod
-    def _cosine_distance(embedding1: List[float], embedding2: List[float]) -> float:
+    def _cosine_distance(embedding1: list[float], embedding2: list[float]) -> float:
         """
         Calculate cosine distance between two embeddings.
         """
@@ -291,7 +291,7 @@ class EmbeddingBasedDocumentSplitter:
         return 1.0 - cosine_sim
 
     @staticmethod
-    def _create_splits_from_points(sentence_groups: List[str], split_points: List[int]) -> List[str]:
+    def _create_splits_from_points(sentence_groups: list[str], split_points: list[int]) -> list[str]:
         """
         Create splits based on split points.
         """
@@ -315,7 +315,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return splits
 
-    def _merge_small_splits(self, splits: List[str]) -> List[str]:
+    def _merge_small_splits(self, splits: list[str]) -> list[str]:
         """
         Merge splits that are below min_length.
         """
@@ -341,7 +341,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return merged
 
-    def _split_large_splits(self, splits: List[str]) -> List[str]:
+    def _split_large_splits(self, splits: list[str]) -> list[str]:
         """
         Recursively split splits that are above max_length.
 
@@ -375,7 +375,7 @@ class EmbeddingBasedDocumentSplitter:
         return final_splits
 
     @staticmethod
-    def _create_documents_from_splits(splits: List[str], original_doc: Document) -> List[Document]:
+    def _create_documents_from_splits(splits: list[str], original_doc: Document) -> list[Document]:
         """
         Create Document objects from splits.
         """
@@ -405,7 +405,7 @@ class EmbeddingBasedDocumentSplitter:
 
         return documents
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
         """
@@ -422,7 +422,7 @@ class EmbeddingBasedDocumentSplitter:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "EmbeddingBasedDocumentSplitter":
+    def from_dict(cls, data: dict[str, Any]) -> "EmbeddingBasedDocumentSplitter":
         """
         Deserializes the component from a dictionary.
         """

haystack_experimental/components/preprocessors/md_header_level_inferrer.py

@@ -24,7 +24,7 @@ class MarkdownHeaderLevelInferrer:
     from haystack_experimental.components.preprocessors import MarkdownHeaderLevelInferrer
 
     # Create a document with uniform header levels
-    text = "## Title\nSome content\n## Section\nMore content\n## Subsection\nFinal content"
+    text = "## Title\n## Subheader\nSection\n## Subheader\nMore Content"
     doc = Document(content=text)
 
     # Initialize the inferrer and process the document
@@ -33,7 +33,7 @@ class MarkdownHeaderLevelInferrer:
 
     # The headers are now normalized with proper hierarchy
     print(result["documents"][0].content)
-    > # Title\nSome content\n## Section\nMore content\n### Subsection\nFinal content
+    > # Title\n## Subheader\nSection\n## Subheader\nMore Content
     ```
     """
 

haystack_experimental/components/retrievers/__init__.py

@@ -3,7 +3,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from haystack_experimental.components.retrievers.chat_message_retriever import ChatMessageRetriever
-from haystack_experimental.components.retrievers.multi_query_embedding_retriever import MultiQueryEmbeddingRetriever
-from haystack_experimental.components.retrievers.multi_query_text_retriever import MultiQueryTextRetriever
 
-_all_ = ["ChatMessageRetriever", "MultiQueryTextRetriever", "MultiQueryEmbeddingRetriever"]
+_all_ = ["ChatMessageRetriever"]

haystack_experimental/components/retrievers/chat_message_retriever.py

@@ -2,11 +2,11 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from haystack import DeserializationError, component, default_from_dict, default_to_dict, logging
 from haystack.core.serialization import import_class_by_name
-from haystack.dataclasses import ChatMessage
+from haystack.dataclasses import ChatMessage, ChatRole
 
 from haystack_experimental.chat_message_stores.types import ChatMessageStore
 
@@ -30,41 +30,40 @@ class ChatMessageRetriever:
     ]
 
     message_store = InMemoryChatMessageStore()
-    message_store.write_messages(messages)
+    message_store.write_messages(chat_history_id="user_456_session_123", messages=messages)
     retriever = ChatMessageRetriever(message_store)
 
-    result = retriever.run()
+    result = retriever.run(chat_history_id="user_456_session_123")
 
     print(result["messages"])
     ```
     """
 
-    def __init__(self, message_store: ChatMessageStore, last_k: int = 10):
+    def __init__(self, chat_message_store: ChatMessageStore, last_k: Optional[int] = 10):
         """
         Create the ChatMessageRetriever component.
 
-        :param message_store:
+        :param chat_message_store:
             An instance of a ChatMessageStore.
         :param last_k:
             The number of last messages to retrieve. Defaults to 10 messages if not specified.
         """
-        self.message_store = message_store
-        if last_k <= 0:
-            raise ValueError(f"last_k must be greater than 0. Currently, the last_k is {last_k}")
+        self.chat_message_store = chat_message_store
+        if last_k and last_k <= 0:
+            raise ValueError(f"last_k must be greater than 0. Currently, last_k is {last_k}")
         self.last_k = last_k
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
         :returns:
             Dictionary with serialized data.
         """
-        message_store = self.message_store.to_dict()
-        return default_to_dict(self, message_store=message_store, last_k=self.last_k)
+        return default_to_dict(self, chat_message_store=self.chat_message_store.to_dict(), last_k=self.last_k)
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "ChatMessageRetriever":
+    def from_dict(cls, data: dict[str, Any]) -> "ChatMessageRetriever":
         """
         Deserializes the component from a dictionary.
 
@@ -74,35 +73,67 @@ class ChatMessageRetriever:
             The deserialized component.
         """
         init_params = data.get("init_parameters", {})
-        if "message_store" not in init_params:
-            raise DeserializationError("Missing 'message_store' in serialization data")
-        if "type" not in init_params["message_store"]:
+        if "chat_message_store" not in init_params:
+            raise DeserializationError("Missing 'chat_message_store' in serialization data")
+        if "type" not in init_params["chat_message_store"]:
             raise DeserializationError("Missing 'type' in message store's serialization data")
 
-        message_store_data = init_params["message_store"]
+        message_store_data = init_params["chat_message_store"]
         try:
             message_store_class = import_class_by_name(message_store_data["type"])
         except ImportError as e:
             raise DeserializationError(f"Class '{message_store_data['type']}' not correctly imported") from e
+        if not hasattr(message_store_class, "from_dict"):
+            raise DeserializationError(f"{message_store_class} does not have from_dict method implemented.")
+        init_params["chat_message_store"] = message_store_class.from_dict(message_store_data)
 
-        data["init_parameters"]["message_store"] = default_from_dict(message_store_class, message_store_data)
         return default_from_dict(cls, data)
 
-    @component.output_types(messages=List[ChatMessage])
-    def run(self, last_k: Optional[int] = None) -> Dict[str, List[ChatMessage]]:
+    @component.output_types(messages=list[ChatMessage])
+    def run(
+        self,
+        chat_history_id: str,
+        *,
+        last_k: Optional[int] = None,
+        current_messages: Optional[list[ChatMessage]] = None,
+    ) -> dict[str, list[ChatMessage]]:
         """
         Run the ChatMessageRetriever
 
+        :param chat_history_id:
+            A unique identifier for the chat session or conversation whose messages should be retrieved.
+            Each `chat_history_id` corresponds to a distinct chat history stored in the underlying ChatMessageStore.
+            For example, use a session ID or conversation ID to isolate messages from different chat sessions.
         :param last_k: The number of last messages to retrieve. This parameter takes precedence over the last_k
             parameter passed to the ChatMessageRetriever constructor. If unspecified, the last_k parameter passed
             to the constructor will be used.
+        :param current_messages:
+            A list of incoming chat messages to combine with the retrieved messages. System messages from this list
+            are prepended before the retrieved history, while all other messages (e.g., user messages) are appended
+            after. This is useful for including new conversational context alongside stored history so the output
+            can be directly used as input to a ChatGenerator or an Agent. If not provided, only the stored messages
+            will be returned.
+
         :returns:
-            - `messages` - The retrieved chat messages.
-        :raises ValueError: If last_k is not None and is less than 1
+            A dictionary with the following key:
+            - `messages` - The retrieved chat messages combined with any provided current messages.
+        :raises ValueError: If last_k is not None and is less than 0.
         """
-        if last_k is not None and last_k <= 0:
-            raise ValueError("last_k must be greater than 0")
+        if last_k is not None and last_k < 0:
+            raise ValueError("last_k must be 0 or greater")
+
+        resolved_last_k = last_k or self.last_k
+        if resolved_last_k == 0:
+            return {"messages": current_messages or []}
+
+        retrieved_messages = self.chat_message_store.retrieve_messages(
+            chat_history_id=chat_history_id, last_k=last_k or self.last_k
+        )
 
-        last_k = last_k or self.last_k
+        if not current_messages:
+            return {"messages": retrieved_messages}
 
-        return {"messages": self.message_store.retrieve()[-last_k:]}
+        # We maintain the order: system messages first, then stored messages, then new user messages
+        system_messages = [msg for msg in current_messages if msg.is_from(ChatRole.SYSTEM)]
+        other_messages = [msg for msg in current_messages if not msg.is_from(ChatRole.SYSTEM)]
+        return {"messages": system_messages + retrieved_messages + other_messages}