Extend llm_context_node definition (#5333)

* Updated node definition

Signed-off-by: Elena Khaustova <ymax70rus@gmail.com>

* Added unit tests

Signed-off-by: Elena Khaustova <ymax70rus@gmail.com>

---------

Signed-off-by: Elena Khaustova <ymax70rus@gmail.com>

Author: ElenaKhaustova

kedro/pipeline/llm_context.py

@@ -14,7 +14,7 @@
 time and automatically assigned readable names based on the returned objects.
 """
 
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass, field
 from typing import Any, NamedTuple, TypeVar
 
@@ -143,14 +143,17 @@ class LLMContextNode(Node):
     ```
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         *,
         outputs: str,
         llm: str,
         prompts: list[str],
         tools: list[_ToolConfig] | None = None,
         name: str | None = None,
+        tags: str | Iterable[str] | None = None,
+        confirms: str | list[str] | None = None,
+        namespace: str | None = None,
     ):
         """Create an LLMContextNode.
 
@@ -160,6 +163,10 @@ def __init__(
             prompts: List of dataset names containing prompt content.
             tools: Optional list of tool configurations created via `tool(...)`.
             name: Optional node name; also used as the logical context identifier.
+            tags: Optional set of tags to be applied to the node.
+            confirms: Optional name or the list of the names of the datasets
+                that should be confirmed.
+            namespace: Optional node namespace.
         """
         inputs = {"llm": llm}
 
@@ -201,18 +208,27 @@ def construct_context(llm: object, **kwargs: dict[str, Any]) -> LLMContext:
 
         # call the Node constructor with the func, inputs, outputs, name
         super().__init__(
-            func=construct_context, inputs=inputs, outputs=outputs, name=name
+            func=construct_context,
+            inputs=inputs,
+            outputs=outputs,
+            name=name,
+            tags=tags,
+            confirms=confirms,
+            namespace=namespace,
         )
 
 
 @experimental
-def llm_context_node(
+def llm_context_node(  # noqa: PLR0913
     *,
     outputs: str,
     llm: str,
     prompts: list[str],
     tools: list[_ToolConfig] | None = None,
     name: str | None = None,
+    tags: str | Iterable[str] | None = None,
+    confirms: str | list[str] | None = None,
+    namespace: str | None = None,
 ) -> Node:
     """
     !!! warning "Experimental"
@@ -229,7 +245,12 @@ def llm_context_node(
         prompts: List of dataset names containing prompt content.
         tools: Optional list of tool configurations created via `tool(...)`.
             Each tool declares the Kedro inputs required to construct it.
-        name: Optional name for the node and for the created context.
+        name: Optional node name; also used as the logical context identifier.
+        tags: Optional set of tags to be applied to the node.
+        confirms: Optional name or the list of the names of the datasets
+            that should be confirmed.
+        namespace: Optional node namespace.
+
 
     Returns:
         A Kedro Node that loads all declared datasets, instantiates tools,
@@ -256,4 +277,7 @@ def llm_context_node(
         prompts=prompts,
         tools=tools,
         name=name,
+        tags=tags,
+        confirms=confirms,
+        namespace=namespace,
     )

tests/pipeline/test_llm_context.py

@@ -2,6 +2,7 @@
 
 from kedro.pipeline.llm_context import (
     LLMContext,
+    LLMContextNode,
     _get_tool_name,
     _normalize_outputs,
     llm_context_node,
@@ -185,3 +186,39 @@ class DummyClass:
 def test_get_tool_name_all_branches(obj, expected_name):
     """_get_tool_name should derive a stable, human-friendly tool name."""
     assert _get_tool_name(obj) == expected_name
+
+
+def test_llm_context_node_sets_tags():
+    """LLMContextNode should propagate tags to the underlying Node."""
+    node_obj = LLMContextNode(
+        outputs="out",
+        llm="llm",
+        prompts=[],
+        tags=["llm", "experimental"],
+    )
+
+    assert node_obj.tags == {"llm", "experimental"}
+
+
+def test_llm_context_node_sets_confirms():
+    """LLMContextNode should propagate confirms to the underlying Node."""
+    node_obj = LLMContextNode(
+        outputs="out",
+        llm="llm",
+        prompts=[],
+        confirms=["llm", "prompt"],
+    )
+
+    assert node_obj.confirms == ["llm", "prompt"]
+
+
+def test_llm_context_node_sets_namespace():
+    """LLMContextNode should propagate namespace to the underlying Node."""
+    node_obj = LLMContextNode(
+        outputs="out",
+        llm="llm",
+        prompts=[],
+        namespace="llm_nodes",
+    )
+
+    assert node_obj.namespace == "llm_nodes"