feat: Expose parameter_binding in @node decorator

Support parameter_binding in @node decorator to allow configuring
dictionary key binding for function nodes. Add tests to verify.

Co-authored-by: Bo Yang <ybo@google.com>
PiperOrigin-RevId: 933952179

Author: boyangsvl

src/google/adk/workflow/_node.py

@@ -19,6 +19,7 @@
 from collections.abc import AsyncGenerator
 from collections.abc import Callable
 from typing import Any
+from typing import Literal
 from typing import overload
 from typing import TYPE_CHECKING
 from typing import TypeVar
@@ -38,7 +39,7 @@
   from ..agents.context import Context
   from ..auth.auth_tool import AuthConfig
 
-T = TypeVar("T", bound=Callable[..., Any])
+T = TypeVar('T', bound=Callable[..., Any])
 
 
 @overload
@@ -50,6 +51,7 @@ def node(
     timeout: float | None = None,
     parallel_worker: bool = False,
     auth_config: AuthConfig | None = None,
+    parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> Callable[
     [T], function_node.FunctionNode | parallel_worker_lib._ParallelWorker
 ]:
@@ -66,6 +68,7 @@ def node(
     timeout: float | None = None,
     parallel_worker: bool = False,
     auth_config: AuthConfig | None = None,
+    parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> base_node.BaseNode:
   ...
 
@@ -79,6 +82,7 @@ def node(
     timeout: float | None = None,
     parallel_worker: bool = False,
     auth_config: AuthConfig | None = None,
+    parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> Any:
   """Decorator or function to wrap a NodeLike in a node or override its properties.
 
@@ -107,6 +111,11 @@ async def my_func3(): ...
     parallel_worker: If True, wraps the node in a _ParallelWorker.
     auth_config: If provided, the framework requests user authentication
       before running the node. Requires rerun_on_resume=True.
+    parameter_binding: How function parameters are bound. ``'state'``
+      (default) binds parameters from ``ctx.state``. ``'node_input'``
+      binds parameters from ``node_input`` dict and infers
+      ``input_schema`` / ``output_schema`` from the function signature
+      (used when the node acts as an agent's tool).
 
   Returns:
     If used as a decorator factory (@node() or @node(...)), returns a decorator.
@@ -126,6 +135,7 @@ def wrapper(
         retry_config=retry_config,
         timeout=timeout,
         auth_config=auth_config,
+        parameter_binding=parameter_binding,
     )
     if parallel_worker:
       return parallel_worker_lib._ParallelWorker(node=built_node)
@@ -142,6 +152,7 @@ def wrapper(
         retry_config=retry_config,
         timeout=timeout,
         auth_config=auth_config,
+        parameter_binding=parameter_binding,
     )
     if parallel_worker:
       return parallel_worker_lib._ParallelWorker(node=built_node)
@@ -166,7 +177,7 @@ def model_post_init(self, __context: Any) -> None:
       # to avoid infinite recursion when its run() method is called.
       # The cloned node preserves the class identity and behavior of the
       # original (essential for LlmAgent and Workflow subclasses).
-      worker_node = self.model_copy(update={"parallel_worker": False})
+      worker_node = self.model_copy(update={'parallel_worker': False})
 
       inner = parallel_worker_lib._ParallelWorker(node=worker_node)
       self._inner_node = inner
@@ -181,7 +192,7 @@ def model_copy(
     copied = super().model_copy(update=update, deep=deep)
 
     if copied.parallel_worker:
-      worker_node = copied.model_copy(update={"parallel_worker": False})
+      worker_node = copied.model_copy(update={'parallel_worker': False})
       copied._inner_node = parallel_worker_lib._ParallelWorker(node=worker_node)
       copied.rerun_on_resume = copied._inner_node.rerun_on_resume
 
@@ -195,7 +206,7 @@ async def run_node_impl(
     Subclasses can directly benefit from advanced flags like parallel_worker
     by providing their custom execution logic here.
     """
-    raise NotImplementedError("run_node_impl must be implemented.")
+    raise NotImplementedError('run_node_impl must be implemented.')
     yield
 
   @override
@@ -205,7 +216,7 @@ async def _run_impl(
     """Dispatches to run_node_impl() or parallel_worker inner node."""
     if self.parallel_worker:
       if self._inner_node is None:
-        raise ValueError("inner_node is not initialized for parallel worker.")
+        raise ValueError('inner_node is not initialized for parallel worker.')
       async for output in self._inner_node.run(ctx=ctx, node_input=node_input):
         yield output
     else:

src/google/adk/workflow/utils/_workflow_graph_utils.py

@@ -18,6 +18,7 @@
 
 from typing import Any
 from typing import cast
+from typing import Literal
 
 from ...tools.base_tool import BaseTool
 from .._base_node import BaseNode
@@ -45,6 +46,7 @@ def build_node(
     retry_config: RetryConfig | None = None,
     timeout: float | None = None,
     auth_config: Any = None,
+    parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> BaseNode:
   """Converts a NodeLike to a BaseNode, wrapping async funcs in FunctionNode.
 
@@ -57,6 +59,11 @@ def build_node(
       wrapped node.
     timeout: If provided, overrides the timeout property of the wrapped node.
     auth_config: If provided, passed to FunctionNode for authentication.
+    parameter_binding: How function parameters are bound. ``'state'``
+      (default) binds parameters from ``ctx.state``. ``'node_input'``
+      binds parameters from ``node_input`` dict and infers
+      ``input_schema`` / ``output_schema`` from the function signature
+      (used when the node acts as an agent's tool).
 
   Returns:
     A BaseNode instance.
@@ -128,6 +135,7 @@ def build_node(
         retry_config=retry_config,
         timeout=timeout,
         auth_config=auth_config,
+        parameter_binding=parameter_binding,
     )
   else:
     raise ValueError(

tests/unittests/workflow/test_workflow_node.py

@@ -160,6 +160,41 @@ def my_func2():
   assert not my_func2.rerun_on_resume
 
 
+def test_node_decorator_parameter_binding():
+  """Tests that @node decorator can configure parameter_binding."""
+
+  @node(parameter_binding="node_input")
+  def my_func(foo: str):
+    return f"Hello {foo}"
+
+  assert isinstance(my_func, FunctionNode)
+  assert my_func.parameter_binding == "node_input"
+
+
+@pytest.mark.asyncio
+async def test_node_decorator_parameter_binding_execution():
+  """Tests execution of a node with parameter_binding='node_input'."""
+
+  async def producer_func() -> dict[str, Any]:
+    return {"foo": "hello", "bar": 42}
+
+  @node(parameter_binding="node_input")
+  def my_func(foo: str, bar: int):
+    return f"{foo}:{bar}"
+
+  wf = Workflow(
+      name="test_agent",
+      edges=[
+          (START, producer_func),
+          (producer_func, my_func),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+
+  by_node = _output_by_node(events)
+  assert ("my_func", "hello:42") in by_node
+
+
 def test_node_function_with_base_node():
   """Tests that node() function returns a copied node when given a BaseNode."""