Skip to content

Comments

feat(go): added experimental ways to define tools, interrupts, and streaming#4797

Open
apascal07 wants to merge 8 commits intoap/go-session-flowfrom
ap/go-x-tools
Open

feat(go): added experimental ways to define tools, interrupts, and streaming#4797
apascal07 wants to merge 8 commits intoap/go-session-flowfrom
ap/go-x-tools

Conversation

@apascal07
Copy link
Collaborator

@apascal07 apascal07 commented Feb 21, 2026

Adds experimental tool definition APIs (ai/x) with simplified function signatures, typed interrupt/resume, and a runtime helper package for use inside tool functions.

These APIs are designed to replace the existing DefineTool and DefineMultipartTool in a future major release. The current implementation wraps the existing tool infrastructure for backwards compatibility, but the intent is to invert this -- the new APIs will become the primary implementation and the old ones will be removed.

Examples

Simple Tools

The new DefineTool (experimentally as DefineXTool) replaces both DefineTool and DefineMultipartTool with a single API. The function receives a plain context.Context instead of ai.ToolContext, making tool functions easier to write and test. Multipart responses are handled via tool.AttachParts rather than requiring a separate function signature:

type WeatherInput struct {
    City string `json:"city" jsonschema:"description=city name"`
}

weatherTool := genkit.DefineXTool(g, "getWeather", "Fetches the weather for a given city",
    func(ctx context.Context, input WeatherInput) (string, error) {
        if input.City == "Paris" {
            return "Sunny, 25°C", nil
        }
        return "Cloudy, 18°C", nil
    },
)

resp, _ := genkit.Generate(ctx, g,
    ai.WithPrompt("What's the weather like in Paris?"),
    ai.WithTools(weatherTool),
)
fmt.Println(resp.Text())

Use tool.AttachParts to return additional content (images, media) without changing the function signature:

genkit.DefineXTool(g, "screenshot", "Takes a screenshot",
    func(ctx context.Context, input ScreenshotInput) (string, error) {
        img := takeScreenshot(input.URL)
        tool.AttachParts(ctx, ai.NewMediaPart("image/png", img))
        return "Screenshot captured", nil
    },
)

Partial Tool Responses (Streaming Progress)

Tools can stream partial responses during execution for client-side display (e.g., progress indicators). Use tool.SendPartial inside any tool function -- it's best-effort and no-ops when streaming isn't available.

Define a status type shared between the tool and client:

type PipelineStatus struct {
    Step     string `json:"step"`
    Progress int    `json:"progress"`
}

Send partial responses from the tool:

processDataTool := aix.NewTool("process_data", "Processes data from a source",
    func(ctx context.Context, input ProcessInput) (string, error) {
        tool.SendPartial(ctx, PipelineStatus{Step: "connecting", Progress: 0})
        connect(input.DataSource)

        tool.SendPartial(ctx, PipelineStatus{Step: "processing", Progress: 50})
        results := process(input.DataSource)

        tool.SendPartial(ctx, PipelineStatus{Step: "done", Progress: 100})
        return fmt.Sprintf("Processed %d records", len(results)), nil
    },
)

On the client side, partial tool responses arrive as ModelResponseChunks during streaming. Use Part.IsPartial() to distinguish progress updates from final tool results:

// GenerateStream
for chunk, _ := range ai.GenerateStream(ctx, r, opts...) {
    for _, p := range chunk.Chunk.ToolResponses() {
        if p.IsPartial() {
            status := p.ToolResponse.Output.(PipelineStatus)
            fmt.Printf("[%s... %d%%]\n", status.Step, status.Progress)
        }
    }
    fmt.Print(sv.Chunk.Text())
}

// Agent flow
for chunk, _ := range conn.Receive() {
    for _, p := range chunk.ModelChunk.ToolResponses() {
        if p.IsPartial() {
            status := p.ToolResponse.Output.(PipelineStatus)
            fmt.Printf("[%s... %d%%]", status.Step, status.Progress)
        }
    }
    fmt.Print(chunk.ModelChunk.Text())
}

Interruptible Tools

DefineInterruptibleTool adds typed interrupt/resume for human-in-the-loop workflows. The function receives a *Res parameter that is non-nil when the tool is being re-executed after an interrupt. Like DefineXTool, multipart responses are supported via tool.AttachParts:

type TransferInput struct {
    ToAccount string  `json:"toAccount"`
    Amount    float64 `json:"amount"`
}

type TransferInterrupt struct {
    Reason string  `json:"reason"`
    Amount float64 `json:"amount"`
}

type Confirmation struct {
    Approved bool `json:"approved"`
}

transferTool := genkit.DefineInterruptibleTool(g, "transfer",
    "Transfers money to another account.",
    func(ctx context.Context, input TransferInput, confirm *Confirmation) (string, error) {
        if confirm != nil && !confirm.Approved {
            return "cancelled", nil
        }
        if confirm == nil && input.Amount > 100 {
            return "", tool.Interrupt(TransferInterrupt{
                Reason: "large_amount",
                Amount: input.Amount,
            })
        }
        return "completed", nil
    },
)

Handle the interrupt on the caller side. Use transferTool.Resume for typed, compile-time-checked resume data:

resp, _ := genkit.Generate(ctx, g,
    ai.WithPrompt("Transfer $200 to Alice"),
    ai.WithTools(transferTool),
)
for _, interrupt := range resp.Interrupts() {
    meta, _ := tool.InterruptAs[TransferInterrupt](interrupt)
    fmt.Printf("Transfer of $%.2f requires confirmation (reason: %s)\n", meta.Amount, meta.Reason)

    // Resume: re-execute the tool with typed resume data.
    restart, _ := transferTool.Resume(interrupt, Confirmation{Approved: true})
    resp, _ = genkit.Generate(ctx, g,
        ai.WithMessages(resp.History()...),
        ai.WithTools(transferTool),
        ai.WithToolRestarts(restart),
    )
}

Or respond with a pre-computed result instead of re-executing the tool:

// Respond: provide the tool's output directly, skipping re-execution.
response, _ := transferTool.Respond(interrupt, "cancelled by user")
resp, _ = genkit.Generate(ctx, g,
    ai.WithMessages(resp.History()...),
    ai.WithTools(transferTool),
    ai.WithToolResponses(response),
)

When the caller does not have access to the tool definition (e.g., in a generic interrupt handler), use the definition-free helpers tool.Resume and tool.Respond:

restart, _ := tool.Resume(interrupt, Confirmation{Approved: true})
// or
response, _ := tool.Respond(interrupt, "cancelled by user")

Agent Flows with Interrupts

Interruptible tools integrate naturally with DefinePromptAgent. The agent streams interrupts to the client, and the client handles user interaction and sends resume data back:

genkit.DefineInterruptibleTool(g, "transferMoney", "Transfers money.", transferFunc)

paymentAgent := genkit.DefinePromptAgent[any, any](g, "paymentAgent", nil)

conn, _ := paymentAgent.StreamBidi(ctx)
conn.SendText("Transfer $200 to Alice")

for chunk, _ := range conn.Receive() {
    fmt.Print(chunk.ModelChunk.Text())
    for _, interrupt := range chunk.ModelChunk.Interrupts() {
        meta, _ := tool.InterruptAs[TransferInterrupt](interrupt)
        fmt.Printf("\n[%s] $%.2f\n", meta.Reason, meta.Amount)

        restart, _ := transferTool.Resume(interrupt, Confirmation{Approved: true})
        conn.SendToolRestarts(restart)
    }
    if chunk.EndTurn {
        break
    }
}

API Reference

Runtime Helpers (ai/x/tool -- experimental)

The tool package provides functions for use inside tool functions:

// Interrupt pauses tool execution and sends data to the caller.
func Interrupt(data any) error

// InterruptAs extracts typed interrupt data from an interrupted tool request part.
func InterruptAs[T any](p *ai.Part) (T, bool)

// Resume creates a restart part for resuming an interrupted tool call.
// Does not require access to the tool definition.
func Resume[Res any](interruptedPart *ai.Part, data Res) (*ai.Part, error)

// Respond creates a tool response part for an interrupted tool request.
// Provides a pre-computed result instead of re-executing the tool.
// Does not require access to the tool definition.
func Respond(interruptedPart *ai.Part, output any) (*ai.Part, error)

// AttachParts attaches additional content parts (e.g., media) to the tool's
// response without changing the function signature.
func AttachParts(ctx context.Context, parts ...*ai.Part)

// OriginalInput returns the original input if the caller replaced it during restart.
func OriginalInput[In any](ctx context.Context) (In, bool)

// SendPartial streams a partial tool response during execution (e.g., progress).
// Best-effort: no-ops when streaming is not available.
func SendPartial(ctx context.Context, output any)

// SendChunk streams a raw ModelResponseChunk during tool execution.
// Unlike SendPartial (which wraps data as a partial tool response),
// this gives the tool full control over the chunk contents.
// Best-effort: no-ops when streaming is not available.
func SendChunk(ctx context.Context, chunk *ai.ModelResponseChunk)

Partial Tool Response Helpers (ai package)

// Part.IsPartial reports whether a Part is a partial (streaming) tool response.
func (p *Part) IsPartial() bool

// NewPartialToolResponsePart creates a Part marked as a partial tool response.
func NewPartialToolResponsePart(r *ToolResponse) *Part

// ModelResponseChunk.ToolResponses returns tool response parts from a chunk.
// Use Part.IsPartial() to distinguish progress updates from final results.
func (c *ModelResponseChunk) ToolResponses() []*Part

Tool Definitions (ai/x -- experimental)

Define & Register

// DefineTool creates a tool with a plain context.Context signature and registers it.
func DefineTool[In, Out any](
    r Registry, name, description string,
    fn ToolFunc[In, Out],
    opts ...ai.ToolOption,
) *Tool[In, Out]

// DefineInterruptibleTool creates a tool with typed interrupt/resume and registers it.
func DefineInterruptibleTool[In, Out, Res any](
    r Registry, name, description string,
    fn InterruptibleToolFunc[In, Out, Res],
    opts ...ai.ToolOption,
) *InterruptibleTool[In, Out, Res]

Create Without Registering

// NewTool creates an unregistered tool.
func NewTool[In, Out any](name, description string, fn ToolFunc[In, Out], opts ...ai.ToolOption) *Tool[In, Out]

// NewInterruptibleTool creates an unregistered interruptible tool.
func NewInterruptibleTool[In, Out, Res any](name, description string, fn InterruptibleToolFunc[In, Out, Res], opts ...ai.ToolOption) *InterruptibleTool[In, Out, Res]

Genkit Convenience Functions

// DefineXTool registers via the Genkit instance.
func DefineXTool[In, Out any](g *Genkit, name, description string, fn aix.ToolFunc[In, Out], opts ...ai.ToolOption) *aix.Tool[In, Out]

// DefineInterruptibleTool registers via the Genkit instance.
func DefineInterruptibleTool[In, Out, Res any](g *Genkit, name, description string, fn aix.InterruptibleToolFunc[In, Out, Res], opts ...ai.ToolOption) *aix.InterruptibleTool[In, Out, Res]

Function Signatures

// ToolFunc -- simple tools.
type ToolFunc[In, Out any] = func(ctx context.Context, input In) (Out, error)

// InterruptibleToolFunc -- tools with typed resume.
// The resumed parameter is non-nil when the tool is re-executed after an interrupt.
type InterruptibleToolFunc[In, Out, Resume any] = func(ctx context.Context, input In, res *Resume) (Out, error)

InterruptibleTool[In, Out, Res]

// Resume creates a restart part for this tool with typed data.
// Validates that the interrupted part belongs to this tool.
func (*InterruptibleTool) Resume(interruptedPart *ai.Part, data Res) (*ai.Part, error)

// Respond creates a tool response part for this tool with typed output.
// Provides a pre-computed result instead of re-executing the tool.
// Validates that the interrupted part belongs to this tool.
func (*InterruptibleTool) Respond(interruptedPart *ai.Part, output Out) (*ai.Part, error)

Tool[In, Out]

Implements ai.Tool -- usable anywhere a tool is accepted (ai.WithTools, etc.).

func (*Tool) Name() string
func (*Tool) Definition() *ai.ToolDefinition
func (*Tool) RunRaw(ctx context.Context, input any) (any, error)
func (*Tool) RunRawMultipart(ctx context.Context, input any) (*ai.MultipartToolResponse, error)
func (*Tool) Respond(toolReq *ai.Part, outputData any, opts *ai.RespondOptions) *ai.Part
func (*Tool) Register(r api.Registry)

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @apascal07, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces experimental features to the Genkit Go SDK, focusing on enhancing the flexibility and control over tool definitions and their execution. It provides developers with more streamlined function signatures for tools, the ability to produce multipart responses, and a robust mechanism for interrupting and resuming tool execution, particularly useful for interactive scenarios requiring user input or confirmation. These changes aim to simplify tool development and enable more complex agent behaviors.

Highlights

  • Experimental Tool Definitions: Introduced new experimental packages go/ai/x/tool and go/ai/x/aix to provide alternative, simplified ways of defining tools and handling tool interruptions in Go.
  • Interruptible Tools: Added support for interruptible tools, allowing tool functions to pause execution and send data to the caller, which can then resume the tool with new data.
  • Simplified Function Signatures: New tool definitions (aix.ToolFunc, aix.InterruptibleToolFunc) now accept a plain context.Context instead of ai.ToolContext, simplifying tool implementation.
  • Multipart Tool Responses: Enabled tools to attach additional content parts (e.g., media) to their responses using tool.AttachParts without altering the function signature.
  • New Sample Application: Included a new sample (x-agent-interrupts) demonstrating how to build an interactive agent that handles tool interrupts for user confirmation or input.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • go/ai/x/tool/tool.go
    • Added new package tool for runtime helpers within tool functions.
    • Introduced InterruptError, Interrupt, InterruptAs, and Resume for managing tool interruptions.
    • Added AttachParts to allow tools to include additional content in their responses.
    • Implemented OriginalInput to retrieve the initial input during tool restarts.
  • go/ai/x/tools.go
    • Added new package aix for experimental tool definitions.
    • Defined ToolFunc and InterruptibleToolFunc for simplified tool function signatures.
    • Introduced Tool and InterruptibleTool structs to wrap and extend ai.ToolDef with experimental features.
    • Provided DefineTool, NewTool, DefineInterruptibleTool, and NewInterruptibleTool for creating and registering these new tool types.
  • go/genkit/genkit.go
    • Added DefineXTool to define tools with simplified function signatures.
    • Added DefineInterruptibleTool to define tools that support typed interrupt/resume behavior.
  • go/samples/x-agent-interrupts/main.go
    • Added new sample application demonstrating the experimental interruptible tool API.
    • Implemented a transferMoney tool that interrupts for insufficient balance or large transfer confirmation.
    • Showcased client-side handling of tool interrupts and resumption with user interaction.
  • go/samples/x-agent-interrupts/prompts/paymentAgent.prompt
    • Added prompt definition for the paymentAgent in the x-agent-interrupts sample.
Activity
  • The pull request introduces new features and a sample, indicating initial development and implementation.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces an excellent experimental feature for defining tools with simplified signatures and interrupt capabilities. The new tool and aix packages are well-designed, and the sample application effectively demonstrates the new functionality. My review includes suggestions to enhance error handling, improve maintainability in the sample code, and correct a documentation example. Overall, this is a valuable addition to the library.

@apascal07 apascal07 requested a review from pavelgj February 21, 2026 02:16
@apascal07 apascal07 changed the title feat(go): added experimental ways to define tools and interrupts feat(go): added experimental ways to define tools, interrupts, and tool streaming Feb 21, 2026
@apascal07 apascal07 changed the title feat(go): added experimental ways to define tools, interrupts, and tool streaming feat(go): added experimental ways to define tools, interrupts, and streaming Feb 21, 2026
@apascal07 apascal07 marked this pull request as ready for review February 24, 2026 18:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant