feat(py/genkit): add multipart tool support (tool.v2)

Add support for multipart tools that can return both structured
output and rich content parts, matching the JS SDK defineTool
with multipart: true.

Changes:
- Add TOOL_V2 enum value to ActionKind
- Add multipart kwarg to @ai.tool() decorator
- When multipart=True, register as tool.v2 with metadata
  type=tool.v2 and tool.multipart=True
- When multipart=False (default), register both tool and tool.v2
  (v2 wraps output in {output: result}) for discoverability
- Update resolve_tool() and resolve_parameters() to look up
  both tool and tool.v2 kinds

Includes 6 new tests covering registration, metadata,
execution, and v2 wrapper behavior.

Author: yesudeep

py/PARITY_AUDIT.md

@@ -220,7 +220,7 @@ All samples except `provider-checks-hello` had `LICENSE` ✅ (now fixed).
 | `defineStreamingFlow` | ✅ (via options) | ✅ `DefineStreamingFlow` | ✅ (via streaming param) | ✅ |
 | `defineTool` | ✅ | ✅ `DefineTool` | ✅ `.tool()` decorator | ✅ |
 | `defineToolWithInputSchema` | — | ✅ `DefineToolWithInputSchema` | — | Go-only |
-| `defineTool({multipart: true})` | ✅ | ✅ `DefineMultipartTool` | ❌ | ❌ Python missing (G18) |
+| `defineTool({multipart: true})` | ✅ | ✅ `DefineMultipartTool` | ✅ `.tool(multipart=True)` | ✅ (PR #4513) |
 | `defineModel` | ✅ | ✅ `DefineModel` | ✅ `define_model` | ✅ |
 | `defineBackgroundModel` | ✅ | ✅ `DefineBackgroundModel` | ✅ `define_background_model` | ✅ |
 | `definePrompt` | ✅ | ✅ `DefinePrompt` | ✅ `define_prompt` | ✅ |
@@ -330,7 +330,7 @@ Python users typically use `httpx` or `requests` directly.
 | Feature | JS | Go | Python | Gap Owner | Priority |
 |---------|:--:|:--:|:------:|-----------|:--------:|
 | `runFlow` / `streamFlow` client | ✅ (beta/client) | ❌ | ❌ | Go + Python | P2 |
-| `defineTool({multipart: true})` | ✅ | ✅ | ❌ | Python | P1 |
+| `defineTool({multipart: true})` | ✅ | ✅ | ✅ | — | ✅ Done (PR #4513) |
 | Model API V2 (`apiVersion: 'v2'`) | ✅ | ❌ | ❌ | Go + Python | P1 |
 | `defineDynamicActionProvider` | ✅ | ❌ | ✅ | Go | P2 |
 | `defineIndexer` | ✅ | ❌ | ✅ | Go | P2 |
@@ -490,11 +490,11 @@ Full plugin list from the repository README (10 plugins, 33 contributors, 54 rel
 | G14 | Python | Implement `validate_support` middleware | §8f | ⬜ |
 | G15 | Python | Implement `download_request_media` middleware | §8f | ⬜ |
 | G16 | Python | Implement `simulate_system_prompt` middleware | §8f | ⬜ |
-| G18 | Python | Add multipart tool support (`defineTool({multipart: true})`) | §8h | ⬜ |
+| G18 | Python | Add multipart tool support (`defineTool({multipart: true})`) | §8h | ✅ PR #4513 |
 | G19 | Python | Add Model API V2 (`defineModel({apiVersion: 'v2'})`) | §8i | ⬜ |
-| G20 | Python | Add `context` parameter to `Genkit()` constructor | §8j | ⬜ |
-| G21 | Python | Add `clientHeader` parameter to `Genkit()` constructor | §8j | ⬜ |
-| G22 | Python | Add `name` parameter to `Genkit()` constructor | §8j | ⬜ |
+| G20 | Python | Add `context` parameter to `Genkit()` constructor | §8j | ✅ PR #4512 |
+| G21 | Python | Add `clientHeader` parameter to `Genkit()` constructor | §8j | ✅ PR #4512 |
+| G22 | Python | Add `name` parameter to `Genkit()` constructor | §8j | ✅ PR #4512 |
 | G4 | Python | Move `augment_with_context` to define-model time | §8b.2 | ⬜ |
 | G9 | Python | Add Pinecone vector store plugin | §5g | ⬜ |
 | G10 | Python | Add ChromaDB vector store plugin | §5g | ⬜ |
@@ -908,10 +908,10 @@ export function apiKey(
 
 | Feature | JS | Python | Gap |
 |---------|:--:|:------:|:---:|
-| `defineTool({multipart: true})` | ✅ Supported. Creates a `MultipartToolAction` of type `tool.v2`. | ❌ Not supported. `define_tool` has no `multipart` parameter. | **G18** |
-| `MultipartToolAction` type | ✅ `tool.ts:107-122` — Action with `tool.v2` type, returns `{output?, content?}`. | ❌ Does not exist. | **G18** |
-| `MultipartToolResponse` type | ✅ `parts.ts` — Schema with `output` and `content` fields. | ⚠️ Type exists in `typing.py:933` but unused in tool definition. | Partial |
-| Auto-registration of `tool.v2` | ✅ Non-multipart tools are also registered as `tool.v2` with wrapped output. | ❌ No dual registration. | **G18** |
+| `defineTool({multipart: true})` | ✅ Supported. Creates a `MultipartToolAction` of type `tool.v2`. | ✅ `.tool(multipart=True)` registers as `tool.v2` with metadata `tool.multipart=True`. | ✅ PR #4513 |
+| `MultipartToolAction` type | ✅ `tool.ts:107-122` — Action with `tool.v2` type, returns `{output?, content?}`. | ✅ Registered under `ActionKind.TOOL_V2` with appropriate metadata. | ✅ PR #4513 |
+| `MultipartToolResponse` type | ✅ `parts.ts` — Schema with `output` and `content` fields. | ✅ Multipart tool functions return `{output?, content?}` dict. | ✅ PR #4513 |
+| Auto-registration of `tool.v2` | ✅ Non-multipart tools are also registered as `tool.v2` with wrapped output. | ✅ Non-multipart tools register both `tool` and `tool.v2` (v2 wraps output in `{output: result}`). | ✅ PR #4513 |
 
 **JS** (`js/ai/src/tool.ts:306-335`):
 ```ts
@@ -1040,11 +1040,11 @@ export interface GenkitOptions {
 | G15 | Python | `download_request_media` middleware missing | P2 | `py/packages/genkit/src/genkit/blocks/middleware.py` | URL media transformed to data URI |
 | G16 | Python | `simulate_system_prompt` missing | P2 | `py/packages/genkit/src/genkit/blocks/middleware.py` | system message rewritten for unsupported model |
 | G17 | Python | `api_key()` context provider missing | P3 | `py/packages/genkit/src/genkit/core/context.py` | auth header extraction + policy callback tests |
-| G18 | Python | multipart tool (`tool.v2`) missing | P1 | `py/packages/genkit/src/genkit/blocks/tools.py`, `.../blocks/generate.py` | tool call returns `output` + `content` parity |
+| G18 | Python | ~~multipart tool (`tool.v2`) missing~~ | P1 | `ai/_registry.py`, `core/action/types.py`, `blocks/generate.py` | ✅ **Done** (PR #4513) |
 | G19 | Python | Model API V2 runner interface missing | P1 | `py/packages/genkit/src/genkit/ai/_registry.py`, `.../blocks/model.py` | v2 model receives unified options struct |
-| G20 | Python | `Genkit(context=...)` missing | P2 | `py/packages/genkit/src/genkit/ai/_aio.py` | context propagates to action executions |
-| G21 | Python | `Genkit(clientHeader=...)` missing | P2 | `py/packages/genkit/src/genkit/ai/_aio.py`, `.../core/http_client.py` | outbound header includes custom token |
-| G22 | Python | `Genkit(name=...)` missing | P2 | `py/packages/genkit/src/genkit/ai/_aio.py`, `.../core/reflection.py` | Dev UI/reflection shows custom name |
+| G20 | Python | ~~`Genkit(context=...)` missing~~ | P2 | `ai/_aio.py`, `core/registry.py` | ✅ **Done** (PR #4512) |
+| G21 | Python | ~~`Genkit(clientHeader=...)` missing~~ | P2 | `ai/_aio.py`, `core/constants.py` | ✅ **Done** (PR #4512) |
+| G22 | Python | ~~`Genkit(name=...)` missing~~ | P2 | `ai/_aio.py`, `ai/_runtime.py`, `core/registry.py` | ✅ **Done** (PR #4512) |
 | G23 | Go | `defineDynamicActionProvider` parity missing | P2 | `go/genkit/genkit.go`, `go/core/registry.go` | DAP action discovery + resolve test |
 | G24 | Go | `defineIndexer` parity missing | P2 | `go/genkit/genkit.go`, `go/ai` indexing action | indexer registration + invoke test |
 | G25 | Go | `defineReranker` + `rerank` runtime missing | P1 | `go/genkit/genkit.go`, `go/ai` reranker block | reranker registration + scoring call |
@@ -1198,9 +1198,9 @@ Reverse topological sort of the gap DAG yields the following dependency levels.
 |----|-----|-----------|----------------|:------:|----------|
 | **P1.1** | **G2** | Add `middleware` storage to `Action` class; implement `action_with_middleware()` wrapper that chains model-level middleware around `action.run()` | `core/action/_action.py` | L | G1, G12, G13, G15, G19 |
 | **P1.2** | **G6** | Update `on_trace_start` callback signature to `(trace_id: str, span_id: str)` throughout action system | `core/action/_action.py`, `core/reflection.py`, `core/trace/` | S | G5 |
-| **P1.3** | **G18** | Add multipart tool support: `define_tool(multipart=True)`, `MultipartToolAction` type `tool.v2`, dual registration for non-multipart tools | `blocks/tools.py`, `blocks/generate.py` | M | — |
-| **P1.4** | **G20** | Add `context` parameter to `Genkit()` that sets `registry.context` for default action context | `ai/_aio.py` | XS | — |
-| **P1.5** | **G21** | Add `clientHeader` parameter to `Genkit()` that appends to `GENKIT_CLIENT_HEADER` via `set_client_header()` | `ai/_aio.py`, `core/http_client.py` | XS | G8 |
+| **P1.3** | **G18** | ~~Add multipart tool support: `define_tool(multipart=True)`, `MultipartToolAction` type `tool.v2`, dual registration for non-multipart tools~~ | `ai/_registry.py`, `core/action/types.py`, `blocks/generate.py` | M | ✅ **Done** (PR #4513) |
+| **P1.4** | **G20** | ~~Add `context` parameter to `Genkit()` that sets `registry.context` for default action context~~ | `ai/_aio.py`, `core/registry.py` | XS | ✅ **Done** (PR #4512) |
+| **P1.5** | **G21** | ~~Add `clientHeader` parameter to `Genkit()` that appends to `GENKIT_CLIENT_HEADER` via `set_client_header()`~~ | `ai/_aio.py`, `core/constants.py` | XS | ✅ **Done** (PR #4512) |
 
 **Exit criteria**: All unit tests green for action middleware dispatch, span_id propagation, tool.v2 registration, and constructor parameter propagation.
 
@@ -1439,8 +1439,8 @@ Milestone     ▲ P1 infra    ▲ Middleware     ▲ Full P1    ▲ Client
 |----|:-----:|------|----------|:----------:|
 | **PR-1a** | Core | G2 | Add `middleware` list to `Action.__init__()`, implement `action_with_middleware()` dispatch wrapper, unit tests for middleware chaining | — |
 | **PR-1b** | Core | G6 | Update `on_trace_start` callback signature to `(trace_id, span_id)` across action system + tracing, update all call sites | — |
-| **PR-1c** | Core | G18 | Multipart tool support: `define_tool(multipart=True)`, `tool.v2` action type, dual registration for non-multipart tools, unit tests | — |
-| **PR-1d** | Core | G20, G21 | `Genkit(context=..., client_header=...)` constructor params — small additive changes, can combine in one PR | — |
+| **PR-1c** | Core | G18 | ~~Multipart tool support: `define_tool(multipart=True)`, `tool.v2` action type, dual registration for non-multipart tools, unit tests~~ | ✅ PR #4513 |
+| **PR-1d** | Core | G20, G21, G22 | ~~`Genkit(context=..., client_header=..., name=...)` constructor params~~ | ✅ PR #4512 |
 
 *PR-1a is the critical-path item. Land it first to unblock Phase 2.*
 

py/packages/genkit/src/genkit/ai/_registry.py

@@ -522,15 +522,21 @@ async def get_tools():
         return define_dap_block(self.registry, config, fn)
 
     def tool(
-        self, name: str | None = None, description: str | None = None
+        self, name: str | None = None, description: str | None = None, *, multipart: bool = False
     ) -> Callable[[Callable[P, T]], Callable[P, T]]:
         """Decorator to register a function as a tool.
 
         Args:
-            name: Optional name for the flow. If not provided, uses the function
+            name: Optional name for the tool. If not provided, uses the function
                 name.
             description: Description for the tool to be passed to the model;
                 if not provided, uses the function docstring.
+            multipart: If True, the tool is registered as a multipart tool
+                (``tool.v2``). The function should return a dict with optional
+                ``output`` and ``content`` keys. If False (default), both a
+                ``tool`` and a ``tool.v2`` wrapper action are registered so that
+                the tool is discoverable under both kinds. Mirrors JS SDK's
+                ``defineTool({ multipart: true })``.
 
         Returns:
             A decorator function that registers the tool.
@@ -564,14 +570,40 @@ def tool_fn_wrapper(*args: Any) -> Any:  # noqa: ANN401
                     case _:
                         raise ValueError('tool must have 0-2 args...')
 
+            tool_kind = cast(ActionKind, ActionKind.TOOL_V2 if multipart else ActionKind.TOOL)
+            tool_metadata: dict[str, object] = {'type': 'tool.v2' if multipart else 'tool'}
+            if multipart:
+                tool_metadata['tool'] = {'multipart': True}
+
             action = self.registry.register_action(
                 name=tool_name,
-                kind=cast(ActionKind, ActionKind.TOOL),
+                kind=tool_kind,
                 description=tool_description,
                 fn=tool_fn_wrapper,
                 metadata_fn=func,
+                metadata=tool_metadata,
             )
 
+            # For non-multipart tools, also register a tool.v2 wrapper that
+            # wraps the output in {output: result} so all tools are
+            # discoverable under tool.v2, matching JS SDK behavior.
+            if not multipart:
+
+                async def v2_wrapper_fn(*args: Any) -> dict[str, object]:  # noqa: ANN401
+                    result = tool_fn_wrapper(*args)
+                    if asyncio.iscoroutine(result):
+                        result = await result
+                    return {'output': result}
+
+                self.registry.register_action(
+                    name=tool_name,
+                    kind=cast(ActionKind, ActionKind.TOOL_V2),
+                    description=tool_description,
+                    fn=v2_wrapper_fn,
+                    metadata_fn=func,
+                    metadata={'type': 'tool.v2'},
+                )
+
             @wraps(func)
             async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:  # noqa: ANN401
                 """Asynchronous wrapper for the tool function.

py/packages/genkit/src/genkit/blocks/generate.py

@@ -626,9 +626,7 @@ async def resolve_parameters(
     tools: list[Action[Any, Any, Any]] = []
     if request.tools:
         for tool_name in request.tools:
-            tool_action = await registry.resolve_action(cast(ActionKind, ActionKind.TOOL), tool_name)
-            if tool_action is None:
-                raise Exception(f'Unable to resolve tool {tool_name}')
+            tool_action = await resolve_tool(registry, tool_name)
             tools.append(tool_action)
 
     format_def: FormatDef | None = None
@@ -842,6 +840,9 @@ async def _resolve_tool_request(tool: Action, tool_request_part: ToolRequestPart
 async def resolve_tool(registry: Registry, tool_name: str) -> Action:
     """Resolve a tool by name from the registry.
 
+    Looks up the tool under both ``tool`` and ``tool.v2`` action kinds,
+    matching the JS SDK's ``lookupToolByName`` behavior.
+
     Args:
         registry: The registry to resolve the tool from.
         tool_name: The name of the tool to resolve.
@@ -853,6 +854,8 @@ async def resolve_tool(registry: Registry, tool_name: str) -> Action:
         ValueError: If the tool could not be resolved.
     """
     tool = await registry.resolve_action(kind=cast(ActionKind, ActionKind.TOOL), name=tool_name)
+    if tool is None:
+        tool = await registry.resolve_action(kind=cast(ActionKind, ActionKind.TOOL_V2), name=tool_name)
     if tool is None:
         raise ValueError(f'Unable to resolve tool {tool_name}')
     return tool

py/packages/genkit/src/genkit/core/action/types.py

@@ -56,6 +56,7 @@ class ActionKind(StrEnum):
     RESOURCE = 'resource'
     RETRIEVER = 'retriever'
     TOOL = 'tool'
+    TOOL_V2 = 'tool.v2'
     UTIL = 'util'
 
 

py/packages/genkit/tests/genkit/ai/multipart_tool_test.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+#
+# Copyright 2025 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for multipart tool support (tool.v2 action kind)."""
+
+from typing import cast
+
+import pytest
+
+from genkit.ai import Genkit
+from genkit.core.action.types import ActionKind
+
+
+@pytest.mark.asyncio
+async def test_regular_tool_registers_both_kinds() -> None:
+    """A non-multipart tool registers under both 'tool' and 'tool.v2'."""
+    ai = Genkit()
+
+    @ai.tool()
+    def add(x: int) -> int:
+        """Add one."""
+        return x + 1
+
+    tool_action = await ai.registry.resolve_action(ActionKind.TOOL, 'add')
+    if tool_action is None:
+        raise AssertionError('Expected tool action registered under ActionKind.TOOL')
+
+    v2_action = await ai.registry.resolve_action(ActionKind.TOOL_V2, 'add')
+    if v2_action is None:
+        raise AssertionError('Expected tool.v2 wrapper action registered under ActionKind.TOOL_V2')
+
+
+@pytest.mark.asyncio
+async def test_regular_tool_v2_wrapper_wraps_output() -> None:
+    """The tool.v2 wrapper for a regular tool wraps output in {output: result}."""
+    ai = Genkit()
+
+    @ai.tool()
+    def double(x: int) -> int:
+        """Double."""
+        return x * 2
+
+    v2_action = await ai.registry.resolve_action(ActionKind.TOOL_V2, 'double')
+    if v2_action is None:
+        raise AssertionError('Expected tool.v2 wrapper')
+
+    result = await v2_action.arun(5)
+    if not isinstance(result.response, dict):
+        raise AssertionError(f'Expected dict response, got {type(result.response).__name__}')
+    if result.response.get('output') != 10:
+        raise AssertionError(f'Expected output=10, got {result.response}')
+
+
+@pytest.mark.asyncio
+async def test_multipart_tool_registers_as_tool_v2() -> None:
+    """A multipart tool is registered under 'tool.v2' only."""
+    ai = Genkit()
+
+    @ai.tool(multipart=True)
+    def rich_tool(query: str) -> dict:
+        """Return rich content."""
+        return {'output': f'result for {query}', 'content': [{'text': 'extra'}]}
+
+    # Should be under tool.v2
+    v2_action = await ai.registry.resolve_action(ActionKind.TOOL_V2, 'rich_tool')
+    if v2_action is None:
+        raise AssertionError('Expected multipart tool registered under ActionKind.TOOL_V2')
+
+    # Should NOT be under tool
+    tool_action = await ai.registry.resolve_action(ActionKind.TOOL, 'rich_tool')
+    if tool_action is not None:
+        raise AssertionError('Multipart tool should NOT be registered under ActionKind.TOOL')
+
+
+@pytest.mark.asyncio
+async def test_multipart_tool_metadata() -> None:
+    """Multipart tool has correct metadata: type='tool.v2' and tool.multipart=True."""
+    ai = Genkit()
+
+    @ai.tool(multipart=True)
+    def my_multipart(x: int) -> dict:
+        """Multipart."""
+        return {'output': x}
+
+    v2_action = await ai.registry.resolve_action(ActionKind.TOOL_V2, 'my_multipart')
+    if v2_action is None:
+        raise AssertionError('Expected multipart tool')
+
+    if v2_action.metadata.get('type') != 'tool.v2':
+        raise AssertionError(f'Expected type="tool.v2", got {v2_action.metadata.get("type")!r}')
+
+    tool_meta = v2_action.metadata.get('tool')
+    if not isinstance(tool_meta, dict):
+        raise AssertionError(f'Expected dict for tool metadata, got {type(tool_meta).__name__}')
+    tool_meta_dict = cast(dict[str, object], tool_meta)
+    if tool_meta_dict.get('multipart') is not True:
+        raise AssertionError(f'Expected tool.multipart=True, got {tool_meta!r}')
+
+
+@pytest.mark.asyncio
+async def test_multipart_tool_execution() -> None:
+    """Multipart tool can be executed and returns the function result directly."""
+    ai = Genkit()
+
+    @ai.tool(multipart=True)
+    def search(query: str) -> dict:
+        """Search."""
+        return {'output': 'found it', 'content': [{'text': f'Details for {query}'}]}
+
+    v2_action = await ai.registry.resolve_action(ActionKind.TOOL_V2, 'search')
+    if v2_action is None:
+        raise AssertionError('Expected multipart tool')
+
+    result = await v2_action.arun('test query')
+    response = result.response
+    if not isinstance(response, dict):
+        raise AssertionError(f'Expected dict, got {type(response).__name__}')
+    if response.get('output') != 'found it':
+        raise AssertionError(f'Expected output="found it", got {response.get("output")!r}')
+
+
+@pytest.mark.asyncio
+async def test_regular_tool_metadata_type() -> None:
+    """Regular (non-multipart) tool has metadata type='tool'."""
+    ai = Genkit()
+
+    @ai.tool()
+    def simple(x: int) -> int:
+        """Simple."""
+        return x
+
+    tool_action = await ai.registry.resolve_action(ActionKind.TOOL, 'simple')
+    if tool_action is None:
+        raise AssertionError('Expected tool action')
+
+    if tool_action.metadata.get('type') != 'tool':
+        raise AssertionError(f'Expected type="tool", got {tool_action.metadata.get("type")!r}')

py/pyproject.toml

@@ -160,6 +160,7 @@ framework-dynamic-tools-demo               = { workspace = true }
 framework-evaluator-demo                   = { workspace = true }
 framework-format-demo                      = { workspace = true }
 framework-middleware-demo                  = { workspace = true }
+framework-multipart-tools                  = { workspace = true }
 framework-prompt-demo                      = { workspace = true }
 framework-realtime-tracing-demo            = { workspace = true }
 framework-restaurant-demo                  = { workspace = true }