chore: Unwraps callables in FunctionNode

FunctionNode now unwraps the callable before checking its type (async/sync, gen/coroutine).

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 929645733

Author: DeanChensj

src/google/adk/workflow/_function_node.py

@@ -212,7 +212,11 @@ def __init__(
           ' The node must rerun after credentials are provided.'
       )
 
-    inferred_name = name or getattr(func, '__name__', None)
+    inferred_name = (
+        name
+        or getattr(func, '__name__', None)
+        or getattr(_unwrap_callable(func), '__name__', None)
+    )
     if not inferred_name:
       raise ValueError(
           'FunctionNode must have a name. If the wrapped callable does not'
@@ -497,9 +501,10 @@ async def _run_impl(
 
     kwargs = self._bind_parameters(ctx, node_input)
 
-    if inspect.isasyncgenfunction(self._func):
+    unwrapped_func = _unwrap_callable(self._func)
+    if inspect.isasyncgenfunction(unwrapped_func):
       items = self._func(**kwargs)
-    elif inspect.isgeneratorfunction(self._func):
+    elif inspect.isgeneratorfunction(unwrapped_func):
       items = _sync_to_async_gen(self._func(**kwargs))
     else:
       items = None
@@ -510,7 +515,7 @@ async def _run_impl(
         if event is not None:
           yield event
     else:
-      if inspect.iscoroutinefunction(self._func):
+      if inspect.iscoroutinefunction(unwrapped_func):
         result = await self._func(**kwargs)
       else:  # Sync function
         result = self._func(**kwargs)

tests/unittests/workflow/test_function_node.py

@@ -1736,3 +1736,49 @@ def __call__(self, z: float) -> None:
     obj = MyCallable()
     hints4 = _get_type_hints_cached(obj)
     assert hints4 == {'z': float, 'return': type(None)}
+
+
+@pytest.mark.asyncio
+async def test_function_node_wrapped_partial(request: pytest.FixtureRequest):
+  """Tests that FunctionNode correctly unwraps functools.partial for async/sync generators and coroutines."""
+  import functools
+
+  async def async_gen_fn(
+      prefix: str, ctx: Context
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(output=f'{prefix} from AsyncGen')
+
+  def sync_gen_fn(prefix: str, ctx: Context) -> Generator[Any, None, None]:
+    yield Event(output=f'{prefix} from SyncGen')
+
+  async def async_fn(prefix: str, ctx: Context) -> str:
+    return f'{prefix} from AsyncCoro'
+
+  p_async_gen = functools.partial(async_gen_fn, 'Hello')
+  p_sync_gen = functools.partial(sync_gen_fn, 'Hello')
+  p_async = functools.partial(async_fn, 'Hello')
+
+  agent = Workflow(
+      name='test_workflow_partial_unwrapping',
+      edges=[
+          (START, p_async_gen),
+          (p_async_gen, p_sync_gen),
+          (p_sync_gen, p_async),
+      ],
+  )
+  events, _, _ = await run_workflow(agent)
+
+  assert simplify_events_with_node(events) == [
+      (
+          'test_workflow_partial_unwrapping@1/async_gen_fn@1',
+          {'output': 'Hello from AsyncGen'},
+      ),
+      (
+          'test_workflow_partial_unwrapping@1/sync_gen_fn@1',
+          {'output': 'Hello from SyncGen'},
+      ),
+      (
+          'test_workflow_partial_unwrapping@1/async_fn@1',
+          {'output': 'Hello from AsyncCoro'},
+      ),
+  ]