Genkit MCP

// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // // SPDX-License-Identifier: Apache-2.0 package ai import ( "context" "encoding/json" "errors" "fmt" "maps" "reflect" "github.com/firebase/genkit/go/core" "github.com/firebase/genkit/go/core/api" "github.com/firebase/genkit/go/internal/base" ) var resumedCtxKey = base.NewContextKey[map[string]any]() var origInputCtxKey = base.NewContextKey[any]() // ToolFunc is the function type for tool implementations. type ToolFunc[In, Out any] = func(ctx *ToolContext, input In) (Out, error) // MultipartToolFunc is the function type for multipart tool implementations. // Unlike regular tools that return just an output value, multipart tools // can return both an output value and additional content parts (like media). type MultipartToolFunc[In any] = func(ctx *ToolContext, input In) (*MultipartToolResponse, error) // ToolRef is a reference to a tool. type ToolRef interface { Name() string } // ToolName is a distinct type for a tool name. // It is meant to be passed where a ToolRef is expected but no Tool is had. type ToolName string // Name returns the name of the tool. func (t ToolName) Name() string { return (string)(t) } // tool is an action with functions specific to tools. // Internally, all tools use the v2 format (returning MultipartToolResponse). // For regular tools, RunRaw unwraps the Output field for backward compatibility. type tool struct { api.Action multipart bool // Whether this is a multipart-only tool. registry api.Registry // Registry for schema resolution. Set when registered. } // Tool represents a tool that can be called by a model. type Tool interface { // Name returns the name of the tool. Name() string // Definition returns the definition for this tool to be passed to models. Definition() *ToolDefinition // RunRaw runs this tool using the provided raw input and returns just the output. RunRaw(ctx context.Context, input any) (any, error) // RunRawMultipart runs this tool and returns the full [MultipartToolResponse]. RunRawMultipart(ctx context.Context, input any) (*MultipartToolResponse, error) // Respond constructs a [Part] with a [ToolResponse] for a given interrupted tool request. Respond(toolReq *Part, outputData any, opts *RespondOptions) *Part // Restart constructs a [Part] with a new [ToolRequest] to re-trigger a tool, // potentially with new input and metadata. Restart(toolReq *Part, opts *RestartOptions) *Part // Register registers the tool with the given registry. Register(r api.Registry) } // toolInterruptError represents an intentional interruption of tool execution. type toolInterruptError struct { Metadata map[string]any } func (e *toolInterruptError) Error() string { return "tool execution interrupted" } // IsToolInterruptError determines whether the error is an interrupt error returned by the tool. func IsToolInterruptError(err error) (bool, map[string]any) { var tie *toolInterruptError if errors.As(err, &tie) { return true, tie.Metadata } return false, nil } // InterruptOptions provides configuration for tool interruption. type InterruptOptions struct { Metadata map[string]any } // RestartOptions provides configuration options for restarting a tool. type RestartOptions struct { // ReplaceInput allows replacing the existing input arguments to the tool with different ones, // for example if the user revised an action before confirming. When input is replaced, // the existing tool request will be amended in the message history. ReplaceInput any // ResumedMetadata is the metadata you want to provide to the tool to aide in reprocessing. // Defaults to true if none is supplied. ResumedMetadata any } // RespondOptions provides configuration options for responding to a tool request. type RespondOptions struct { // Metadata is additional metadata to include in the response. Metadata map[string]any } // ToolContext provides context and utility functions for tool execution. type ToolContext struct { context.Context // Interrupt is a function that can be used to interrupt the tool execution. // Interrupting tool execution returns the control to the caller with the // total model response so far. Interrupt func(opts *InterruptOptions) error // Resumed is optional metadata that can be used to resume the tool execution. // Map is not nil only if the tool was interrupted. Resumed map[string]any // OriginalInput is the original input to the tool if the tool was interrupted, otherwise nil. OriginalInput any } // DefineTool creates a new [Tool] and registers it. // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. func DefineTool[In, Out any]( r api.Registry, name, description string, fn ToolFunc[In, Out], opts ...ToolOption, ) Tool { toolOpts := &toolOptions{} for _, opt := range opts { if err := opt.applyTool(toolOpts); err != nil { panic(fmt.Errorf("ai.DefineTool %q: %w", name, err)) } } // If the user provided a custom input schema, enforce that In is 'any' if toolOpts.InputSchema != nil { typ := reflect.TypeFor[*In]() if typ != nil && typ.Elem().Kind() != reflect.Interface { panic(fmt.Errorf("ai.DefineTool %q: WithInputSchema requires In to be of type 'any', but got %v", name, typ.Elem())) } } metadata, wrappedFn := wrapToolFunc(name, description, fn) action := core.DefineAction(r, name, api.ActionTypeToolV2, metadata, toolOpts.InputSchema, wrappedFn) // Also register under the "tool" action type for backward compatibility. provider, id := api.ParseName(name) r.RegisterAction(api.NewKey(api.ActionTypeTool, provider, id), action) return &tool{Action: action, multipart: false, registry: r} } // DefineToolWithInputSchema creates a new [Tool] with a custom input schema and registers it. // // Deprecated: Use [DefineTool] with [WithInputSchema] instead. func DefineToolWithInputSchema[Out any]( r api.Registry, name, description string, inputSchema map[string]any, fn ToolFunc[any, Out], ) Tool { return DefineTool(r, name, description, fn, WithInputSchema(inputSchema)) } // NewTool creates a new [Tool]. It can be passed directly to [Generate]. // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. func NewTool[In, Out any](name, description string, fn ToolFunc[In, Out], opts ...ToolOption) Tool { toolOpts := &toolOptions{} for _, opt := range opts { if err := opt.applyTool(toolOpts); err != nil { panic(fmt.Errorf("ai.NewTool %q: %w", name, err)) } } // If the user provided a custom input schema, enforce that In is 'any' if toolOpts.InputSchema != nil { var zeroIn *In typ := reflect.TypeOf(zeroIn) if typ != nil && typ.Elem().Kind() != reflect.Interface { panic(fmt.Errorf("ai.NewTool %q: WithInputSchema requires In to be of type 'any', but got %v", name, typ.Elem())) } } metadata, wrappedFn := wrapToolFunc(name, description, fn) metadata["dynamic"] = true action := core.NewAction(name, api.ActionTypeToolV2, metadata, toolOpts.InputSchema, wrappedFn) return &tool{Action: action, multipart: false} } // NewToolWithInputSchema creates a new [Tool] with a custom input schema. It can be passed directly to [Generate]. // // Deprecated: Use [NewTool] with [WithInputSchema] instead. func NewToolWithInputSchema[Out any](name, description string, inputSchema map[string]any, fn ToolFunc[any, Out]) Tool { return NewTool(name, description, fn, WithInputSchema(inputSchema)) } // DefineMultipartTool creates a new multipart [Tool] and registers it. // Multipart tools can return both output data and additional content parts (like media). // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. func DefineMultipartTool[In any]( r api.Registry, name, description string, fn MultipartToolFunc[In], opts ...ToolOption, ) Tool { toolOpts := &toolOptions{} for _, opt := range opts { if err := opt.applyTool(toolOpts); err != nil { panic(fmt.Errorf("ai.DefineMultipartTool %q: %w", name, err)) } } metadata, wrappedFn := wrapMultipartToolFunc(name, description, fn) action := core.DefineAction(r, name, api.ActionTypeToolV2, metadata, toolOpts.InputSchema, wrappedFn) return &tool{Action: action, multipart: true, registry: r} } // NewMultipartTool creates a new multipart [Tool]. It can be passed directly to [Generate]. // Multipart tools can return both output data and additional content parts (like media). // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. func NewMultipartTool[In any](name, description string, fn MultipartToolFunc[In], opts ...ToolOption) Tool { toolOpts := &toolOptions{} for _, opt := range opts { if err := opt.applyTool(toolOpts); err != nil { panic(fmt.Errorf("ai.NewMultipartTool %q: %w", name, err)) } } metadata, wrappedFn := wrapMultipartToolFunc(name, description, fn) metadata["dynamic"] = true action := core.NewAction(name, api.ActionTypeToolV2, metadata, toolOpts.InputSchema, wrappedFn) return &tool{Action: action, multipart: true} } // wrapToolFunc wraps a regular tool function to return MultipartToolResponse. func wrapToolFunc[In, Out any](name, description string, fn ToolFunc[In, Out]) (map[string]any, func(context.Context, In) (*MultipartToolResponse, error)) { var o Out var originalOutputSchema map[string]any if reflect.TypeOf(o) != nil { originalOutputSchema = core.InferSchemaMap(o) } metadata := map[string]any{ "type": api.ActionTypeToolV2, "name": name, "description": description, "tool": map[string]any{"multipart": false}, } if originalOutputSchema != nil { metadata["originalOutputSchema"] = originalOutputSchema } wrappedFn := func(ctx context.Context, input In) (*MultipartToolResponse, error) { toolCtx := &ToolContext{ Context: ctx, Interrupt: func(opts *InterruptOptions) error { return &toolInterruptError{ Metadata: opts.Metadata, } }, Resumed: resumedCtxKey.FromContext(ctx), OriginalInput: origInputCtxKey.FromContext(ctx), } output, err := fn(toolCtx, input) if err != nil { return nil, err } return &MultipartToolResponse{Output: output}, nil } return metadata, wrappedFn } // wrapMultipartToolFunc wraps a multipart tool function. func wrapMultipartToolFunc[In any](name, description string, fn MultipartToolFunc[In]) (map[string]any, func(context.Context, In) (*MultipartToolResponse, error)) { metadata := map[string]any{ "type": api.ActionTypeToolV2, "name": name, "description": description, "tool": map[string]any{"multipart": true}, } wrappedFn := func(ctx context.Context, input In) (*MultipartToolResponse, error) { toolCtx := &ToolContext{ Context: ctx, Interrupt: func(opts *InterruptOptions) error { return &toolInterruptError{ Metadata: opts.Metadata, } }, Resumed: resumedCtxKey.FromContext(ctx), OriginalInput: origInputCtxKey.FromContext(ctx), } return fn(toolCtx, input) } return metadata, wrappedFn } // Definition returns [ToolDefinition] for for this tool. func (t *tool) Definition() *ToolDefinition { desc := t.Action.Desc() // Resolve the input schema if it contains a $ref. inputSchema := desc.InputSchema if t.registry != nil { if resolved, err := core.ResolveSchema(t.registry, inputSchema); err == nil { inputSchema = resolved } } // Use the original output schema if available (for non-multipart tools). outputSchema := desc.OutputSchema if origSchema, ok := desc.Metadata["originalOutputSchema"].(map[string]any); ok { outputSchema = origSchema } // Resolve the output schema if it contains a $ref. if t.registry != nil { if resolved, err := core.ResolveSchema(t.registry, outputSchema); err == nil { outputSchema = resolved } } return &ToolDefinition{ Name: desc.Name, Description: desc.Description, InputSchema: inputSchema, OutputSchema: outputSchema, Metadata: map[string]any{ "multipart": t.multipart, }, } } // Register registers the tool with the given registry. func (t *tool) Register(r api.Registry) { t.registry = r t.Action.Register(r) if !t.multipart { // Also register under the "tool" key for backward compatibility. provider, id := api.ParseName(t.Action.Name()) r.RegisterAction(api.NewKey(api.ActionTypeTool, provider, id), t.Action) } } // RunRaw runs this tool using the provided raw map format data (JSON parsed as map[string]any). func (t *tool) RunRaw(ctx context.Context, input any) (any, error) { resp, err := t.RunRawMultipart(ctx, input) if err != nil { return nil, err } return resp.Output, nil } // RunRawMultipart runs this tool using the provided raw map format data (JSON parsed as map[string]any). // It returns the full multipart response. func (t *tool) RunRawMultipart(ctx context.Context, input any) (*MultipartToolResponse, error) { if t == nil { return nil, core.NewError(core.INVALID_ARGUMENT, "ai.Tool.RunRawMultipart: tool called on a nil tool; check that all tools are defined") } mi, err := json.Marshal(input) if err != nil { return nil, fmt.Errorf("error marshalling tool input for %v: %v", t.Name(), err) } output, err := t.Action.RunJSON(ctx, mi, nil) if err != nil { return nil, fmt.Errorf("error calling tool %v: %w", t.Name(), err) } var resp MultipartToolResponse if err := json.Unmarshal(output, &resp); err != nil { return nil, fmt.Errorf("error parsing tool output for %v: %v", t.Name(), err) } return &resp, nil } // LookupTool looks up the tool in the registry by provided name and returns it. // It checks for "tool.v2" first, then falls back to "tool" for legacy compatibility. func LookupTool(r api.Registry, name string) Tool { if name == "" { return nil } provider, id := api.ParseName(name) // First try tool.v2 (all new tools are registered here) key := api.NewKey(api.ActionTypeToolV2, provider, id) action := r.ResolveAction(key) // Fall back to tool for legacy compatibility if action == nil { key = api.NewKey(api.ActionTypeTool, provider, id) action = r.ResolveAction(key) } if action == nil { return nil } // Check if it's a multipart-only tool desc := action.Desc() multipart := false if toolMeta, ok := desc.Metadata["tool"].(map[string]any); ok { if mp, ok := toolMeta["multipart"].(bool); ok { multipart = mp } } return &tool{Action: action, multipart: multipart, registry: r} } // IsMultipart returns true if the tool is a multipart tool (tool.v2 only). func (t *tool) IsMultipart() bool { return t.multipart } // Respond creates a tool response for an interrupted tool call to pass to the [WithToolResponses] option to [Generate]. // If the part provided is not a tool request, it returns nil. func (t *tool) Respond(toolReq *Part, output any, opts *RespondOptions) *Part { if toolReq == nil || !toolReq.IsToolRequest() { return nil } if opts == nil { opts = &RespondOptions{} } newToolResp := NewResponseForToolRequest(toolReq, output) newToolResp.Metadata = map[string]any{ "interruptResponse": true, } if opts.Metadata != nil { newToolResp.Metadata["interruptResponse"] = opts.Metadata } return newToolResp } // Restart creates a tool request for an interrupted tool call to pass to the [WithToolRestarts] option to [Generate]. // If the part provided is not a tool request, it returns nil. func (t *tool) Restart(p *Part, opts *RestartOptions) *Part { if p == nil || !p.IsToolRequest() { return nil } if opts == nil { opts = &RestartOptions{} } newInput := p.ToolRequest.Input var originalInput any if opts.ReplaceInput != nil { originalInput = newInput newInput = opts.ReplaceInput } newMeta := maps.Clone(p.Metadata) if newMeta == nil { newMeta = make(map[string]any) } newMeta["resumed"] = true if opts.ResumedMetadata != nil { newMeta["resumed"] = opts.ResumedMetadata } if originalInput != nil { newMeta["replacedInput"] = originalInput } delete(newMeta, "interrupt") newToolReq := NewToolRequestPart(&ToolRequest{ Name: p.ToolRequest.Name, Ref: p.ToolRequest.Ref, Input: newInput, }) newToolReq.Metadata = newMeta return newToolReq } // resolveUniqueTools resolves the list of tool refs to a list of all tool names and new tools that must be registered. // Returns an error if there are tool refs with duplicate names. func resolveUniqueTools(r api.Registry, toolRefs []ToolRef) (toolNames []string, newTools []Tool, err error) { toolMap := make(map[string]bool) for _, toolRef := range toolRefs { name := toolRef.Name() if toolMap[name] { return nil, nil, core.NewError(core.INVALID_ARGUMENT, "duplicate tool %q", name) } toolMap[name] = true toolNames = append(toolNames, name) if LookupTool(r, name) == nil { if tool, ok := toolRef.(Tool); ok { newTools = append(newTools, tool) } } } return toolNames, newTools, nil }