option.go•30.3 kB
// 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"
"github.com/firebase/genkit/go/core"
"github.com/firebase/genkit/go/internal/base"
)
// PromptFn is a function that generates a prompt.
type PromptFn = func(context.Context, any) (string, error)
// MessagesFn is a function that generates messages.
type MessagesFn = func(context.Context, any) ([]*Message, error)
// configOptions holds configuration options.
type configOptions struct {
Config any // Primitive (model, embedder, retriever, etc) configuration.
}
// ConfigOption is an option for model configuration.
type ConfigOption interface {
applyConfig(*configOptions) error
applyCommonGen(*commonGenOptions) error
applyPrompt(*promptOptions) error
applyGenerate(*generateOptions) error
applyPromptExecute(*promptExecutionOptions) error
applyEmbedder(*embedderOptions) error
applyRetriever(*retrieverOptions) error
applyEvaluator(*evaluatorOptions) error
}
// applyConfig applies the option to the config options.
func (o *configOptions) applyConfig(opts *configOptions) error {
if o.Config != nil {
if opts.Config != nil {
return errors.New("cannot set config more than once (WithConfig)")
}
opts.Config = o.Config
}
return nil
}
// applyCommonGen applies the option to the common options.
func (o *configOptions) applyCommonGen(opts *commonGenOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyPrompt applies the option to the prompt options.
func (o *configOptions) applyPrompt(opts *promptOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyGenerate applies the option to the generate options.
func (o *configOptions) applyGenerate(opts *generateOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyPromptExecute applies the option to the prompt generate options.
func (o *configOptions) applyPromptExecute(opts *promptExecutionOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyEmbedder applies the option to the embed options.
func (o *configOptions) applyEmbedder(opts *embedderOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyRetriever applies the option to the retrieve options.
func (o *configOptions) applyRetriever(opts *retrieverOptions) error {
return o.applyConfig(&opts.configOptions)
}
// applyEvaluator applies the option to the evaluate options.
func (o *configOptions) applyEvaluator(opts *evaluatorOptions) error {
return o.applyConfig(&opts.configOptions)
}
// WithConfig sets the configuration.
func WithConfig(config any) ConfigOption {
return &configOptions{Config: config}
}
// commonGenOptions are common options for model generation, prompt definition, and prompt execution.
type commonGenOptions struct {
configOptions
Model ModelArg // Model to use.
MessagesFn MessagesFn // Function to generate messages.
Tools []ToolRef // References to tools to use.
Resources []Resource // Resources to be temporarily available during generation.
ToolChoice ToolChoice // Whether tool calls are required, disabled, or optional.
MaxTurns int // Maximum number of tool call iterations.
ReturnToolRequests *bool // Whether to return tool requests instead of making the tool calls and continuing the generation.
Middleware []ModelMiddleware // Middleware to apply to the model request and model response.
}
type CommonGenOption interface {
applyCommonGen(*commonGenOptions) error
applyPrompt(*promptOptions) error
applyGenerate(*generateOptions) error
applyPromptExecute(*promptExecutionOptions) error
}
// applyCommonGen applies the option to the common options.
func (o *commonGenOptions) applyCommonGen(opts *commonGenOptions) error {
if err := o.configOptions.applyConfig(&opts.configOptions); err != nil {
return err
}
if o.MessagesFn != nil {
if opts.MessagesFn != nil {
return errors.New("cannot set messages more than once (either WithMessages or WithMessagesFn)")
}
opts.MessagesFn = o.MessagesFn
}
if o.Model != nil {
if opts.Model != nil {
return errors.New("cannot set model more than once (either WithModel or WithModelName)")
}
opts.Model = o.Model
}
if o.Tools != nil {
if opts.Tools != nil {
return errors.New("cannot set tools more than once (WithTools)")
}
opts.Tools = o.Tools
}
if o.Resources != nil {
if opts.Resources != nil {
return errors.New("cannot set resources more than once (WithResources)")
}
opts.Resources = o.Resources
}
if o.ToolChoice != "" {
if opts.ToolChoice != "" {
return errors.New("cannot set tool choice more than once (WithToolChoice)")
}
opts.ToolChoice = o.ToolChoice
}
if o.MaxTurns > 0 {
if opts.MaxTurns > 0 {
return errors.New("cannot set max turns more than once (WithMaxTurns)")
}
opts.MaxTurns = o.MaxTurns
}
if o.ReturnToolRequests != nil {
if opts.ReturnToolRequests != nil {
return errors.New("cannot configure returning tool requests more than once (WithReturnToolRequests)")
}
opts.ReturnToolRequests = o.ReturnToolRequests
}
if o.Middleware != nil {
if opts.Middleware != nil {
return errors.New("cannot set middleware more than once (WithMiddleware)")
}
opts.Middleware = o.Middleware
}
return nil
}
// applyPromptExecute applies the option to the prompt request options.
func (o *commonGenOptions) applyPromptExecute(reqOpts *promptExecutionOptions) error {
return o.applyCommonGen(&reqOpts.commonGenOptions)
}
// applyPrompt applies the option to the prompt options.
func (o *commonGenOptions) applyPrompt(pOpts *promptOptions) error {
return o.applyCommonGen(&pOpts.commonGenOptions)
}
// applyGenerate applies the option to the generate options.
func (o *commonGenOptions) applyGenerate(genOpts *generateOptions) error {
return o.applyCommonGen(&genOpts.commonGenOptions)
}
// WithMessages sets the messages.
// These messages will be sandwiched between the system and user prompts.
func WithMessages(messages ...*Message) CommonGenOption {
return &commonGenOptions{
MessagesFn: func(ctx context.Context, _ any) ([]*Message, error) {
return messages, nil
},
}
}
// WithMessagesFn sets the request messages to the result of the function.
// These messages will be sandwiched between the system and user messages.
func WithMessagesFn(fn MessagesFn) CommonGenOption {
return &commonGenOptions{MessagesFn: fn}
}
// WithTools sets the tools to use for the generate request.
func WithTools(tools ...ToolRef) CommonGenOption {
return &commonGenOptions{Tools: tools}
}
// WithModel sets either a [Model] or a [ModelRef] that may contain a config.
// Passing [WithConfig] will take precedence over the config in WithModel.
func WithModel(model ModelArg) CommonGenOption {
return &commonGenOptions{Model: model}
}
// WithModelName sets the model name to call for generation.
// The model name will be resolved to a [Model] and may error if the reference is invalid.
func WithModelName(name string) CommonGenOption {
return &commonGenOptions{Model: NewModelRef(name, nil)}
}
// WithMiddleware sets middleware to apply to the model request.
func WithMiddleware(middleware ...ModelMiddleware) CommonGenOption {
return &commonGenOptions{Middleware: middleware}
}
// WithMaxTurns sets the maximum number of tool call iterations before erroring.
// A tool call happens when tools are provided in the request and a model decides to call one or more as a response.
// Each round trip, including multiple tools in parallel, counts as one turn.
func WithMaxTurns(maxTurns int) CommonGenOption {
return &commonGenOptions{MaxTurns: maxTurns}
}
// WithReturnToolRequests configures whether to return tool requests instead of making the tool calls and continuing the generation.
func WithReturnToolRequests(returnReqs bool) CommonGenOption {
return &commonGenOptions{ReturnToolRequests: &returnReqs}
}
// WithToolChoice configures whether by default tool calls are required, disabled, or optional for the prompt.
func WithToolChoice(toolChoice ToolChoice) CommonGenOption {
return &commonGenOptions{ToolChoice: toolChoice}
}
// promptOptions are options for defining a prompt.
type promptOptions struct {
commonGenOptions
promptingOptions
outputOptions
Description string // Description of the prompt.
InputSchema map[string]any // Schema of the input.
DefaultInput map[string]any // Default input that will be used if no input is provided.
Metadata map[string]any // Arbitrary metadata.
}
// PromptOption is an option for defining a prompt.
// It applies only to DefinePrompt().
type PromptOption interface {
applyPrompt(*promptOptions) error
}
// applyPrompt applies the option to the prompt options.
func (o *promptOptions) applyPrompt(opts *promptOptions) error {
if err := o.commonGenOptions.applyPrompt(opts); err != nil {
return err
}
if err := o.promptingOptions.applyPrompt(opts); err != nil {
return err
}
if err := o.outputOptions.applyPrompt(opts); err != nil {
return err
}
if o.Description != "" {
if opts.Description != "" {
return errors.New("cannot set description more than once (WithDescription)")
}
opts.Description = o.Description
}
if o.InputSchema != nil {
if opts.InputSchema != nil {
return errors.New("cannot set input schema more than once (WithInputType)")
}
opts.InputSchema = o.InputSchema
}
if o.DefaultInput != nil {
if opts.DefaultInput != nil {
return errors.New("cannot set default input more than once (WithInputType)")
}
opts.DefaultInput = o.DefaultInput
}
if o.Metadata != nil {
if opts.Metadata != nil {
return errors.New("cannot set metadata more than once (WithMetadata)")
}
opts.Metadata = o.Metadata
}
return nil
}
// WithDescription sets the description of the prompt.
func WithDescription(description string) PromptOption {
return &promptOptions{Description: description}
}
// WithMetadata sets arbitrary metadata for the prompt.
func WithMetadata(metadata map[string]any) PromptOption {
return &promptOptions{Metadata: metadata}
}
// WithInputType uses the type provided to derive the input schema.
// The inputted value will serve as the default input if no input is given at generation time.
// Only supports structs and map[string]any api.
func WithInputType(input any) PromptOption {
var defaultInput map[string]any
switch v := input.(type) {
case map[string]any:
defaultInput = v
default:
data, err := json.Marshal(input)
if err != nil {
panic(fmt.Errorf("failed to marshal default input (WithInputType): %w", err))
}
err = json.Unmarshal(data, &defaultInput)
if err != nil {
panic(fmt.Errorf("type %T is not supported, only structs and map[string]any are supported (WithInputType)", input))
}
}
return &promptOptions{
InputSchema: core.InferSchemaMap(input),
DefaultInput: defaultInput,
}
}
// promptingOptions are options for the system and user prompts of a prompt or generate request.
type promptingOptions struct {
SystemFn PromptFn // Function to generate the system prompt.
PromptFn PromptFn // Function to generate the user prompt.
}
// PromptingOption is an option for the system and user prompts of a prompt or generate request.
// It applies only to DefinePrompt() and Generate().
type PromptingOption interface {
applyPrompting(*promptingOptions) error
applyPrompt(*promptOptions) error
applyGenerate(*generateOptions) error
}
// applyPrompting applies the option to the prompting options.
func (o *promptingOptions) applyPrompting(opts *promptingOptions) error {
if o.SystemFn != nil {
if opts.SystemFn != nil {
return errors.New("cannot set system text more than once (either WithSystem or WithSystemFn)")
}
opts.SystemFn = o.SystemFn
}
if o.PromptFn != nil {
if opts.PromptFn != nil {
return errors.New("cannot set prompt text more than once (either WithPrompt or WithPromptFn)")
}
opts.PromptFn = o.PromptFn
}
return nil
}
// applyPrompt applies the option to the prompt options.
func (o *promptingOptions) applyPrompt(opts *promptOptions) error {
return o.applyPrompting(&opts.promptingOptions)
}
// applyGenerate applies the option to the generate options.
func (o *promptingOptions) applyGenerate(opts *generateOptions) error {
return o.applyPrompting(&opts.promptingOptions)
}
// WithSystem sets the system prompt message.
// The system prompt is always the first message in the list.
func WithSystem(text string, args ...any) PromptingOption {
return &promptingOptions{
SystemFn: func(ctx context.Context, _ any) (string, error) {
// Avoids a compile-time warning about non-constant text.
t := text
return fmt.Sprintf(t, args...), nil
},
}
}
// WithSystemFn sets the function that generates the system prompt message.
// The system prompt is always the first message in the list.
func WithSystemFn(fn PromptFn) PromptingOption {
return &promptingOptions{SystemFn: fn}
}
// WithPrompt sets the user prompt message.
// The user prompt is always the last message in the list.
func WithPrompt(text string, args ...any) PromptingOption {
return &promptingOptions{
PromptFn: func(ctx context.Context, _ any) (string, error) {
// Avoids a compile-time warning about non-constant text.
t := text
return fmt.Sprintf(t, args...), nil
},
}
}
// WithPromptFn sets the function that generates the user prompt message.
// The user prompt is always the last message in the list.
func WithPromptFn(fn PromptFn) PromptingOption {
return &promptingOptions{PromptFn: fn}
}
// outputOptions are options for the output of a prompt or generate request.
type outputOptions struct {
OutputSchema map[string]any // JSON schema of the output.
OutputFormat string // Format of the output. If OutputSchema is set, this is set to OutputFormatJSON.
OutputInstructions *string // Instructions to add to conform the output to a schema. If nil, default instructions will be added. If empty string, no instructions will be added.
CustomConstrained bool // Whether generation should use custom constrained output instead of native model constrained output.
}
// OutputOption is an option for the output of a prompt or generate request.
// It applies only to DefinePrompt() and Generate().
type OutputOption interface {
applyOutput(*outputOptions) error
applyPrompt(*promptOptions) error
applyGenerate(*generateOptions) error
}
// applyOutput applies the option to the output options.
func (o *outputOptions) applyOutput(opts *outputOptions) error {
if o.OutputSchema != nil {
if opts.OutputSchema != nil {
return errors.New("cannot set output schema more than once (WithOutputType)")
}
opts.OutputSchema = o.OutputSchema
}
if o.OutputInstructions != nil {
if opts.OutputInstructions != nil {
return errors.New("cannot set output instructions more than once (WithOutputFormat)")
}
opts.OutputInstructions = o.OutputInstructions
}
if o.OutputFormat != "" {
opts.OutputFormat = o.OutputFormat
}
if o.CustomConstrained {
opts.CustomConstrained = o.CustomConstrained
}
return nil
}
// applyPrompt applies the option to the prompt options.
func (o *outputOptions) applyPrompt(pOpts *promptOptions) error {
return o.applyOutput(&pOpts.outputOptions)
}
// applyGenerate applies the option to the generate options.
func (o *outputOptions) applyGenerate(genOpts *generateOptions) error {
return o.applyOutput(&genOpts.outputOptions)
}
// WithOutputType sets the schema and format of the output based on the value provided.
func WithOutputType(output any) OutputOption {
return &outputOptions{
OutputSchema: base.SchemaAsMap(base.InferJSONSchema(output)),
OutputFormat: OutputFormatJSON,
}
}
// WithOutputFormat sets the format of the output.
func WithOutputFormat(format string) OutputOption {
return &outputOptions{OutputFormat: format}
}
// WithOutputInstructions sets custom instructions for constraining output format in the prompt.
//
// When [WithOutputType] is used without this option, default instructions will be automatically set.
// If you provide empty instructions, no instructions will be added to the prompt.
//
// This will automatically set [WithCustomConstrainedOutput].
func WithOutputInstructions(instructions string) OutputOption {
return &outputOptions{
OutputInstructions: &instructions,
CustomConstrained: true,
}
}
// WithCustomConstrainedOutput opts out of using the model's native constrained output generation.
//
// By default, the system will use the model's native constrained output capabilities when available.
// When this option is set, or when the model doesn't support native constraints, the system will
// use custom implementation to guide the model toward producing properly formatted output.
func WithCustomConstrainedOutput() OutputOption {
return &outputOptions{CustomConstrained: true}
}
// executionOptions are options for the execution of a prompt or generate request.
type executionOptions struct {
Stream ModelStreamCallback // Function to call with each chunk of the generated response.
}
// ExecutionOption is an option for the execution of a prompt or generate request. It applies only to Generate() and prompt.Execute().
type ExecutionOption interface {
applyExecution(*executionOptions) error
applyGenerate(*generateOptions) error
applyPromptExecute(*promptExecutionOptions) error
}
// applyExecution applies the option to the runtime options.
func (o *executionOptions) applyExecution(execOpts *executionOptions) error {
if o.Stream != nil {
if execOpts.Stream != nil {
return errors.New("cannot set stream callback more than once (WithStream)")
}
execOpts.Stream = o.Stream
}
return nil
}
// applyGenerate applies the option to the generate options.
func (o *executionOptions) applyGenerate(genOpts *generateOptions) error {
return o.applyExecution(&genOpts.executionOptions)
}
// applyPromptExecute applies the option to the prompt request options.
func (o *executionOptions) applyPromptExecute(pgOpts *promptExecutionOptions) error {
return o.applyExecution(&pgOpts.executionOptions)
}
// WithStreaming sets the stream callback for the generate request.
// A callback is a function that is called with each chunk of the generated response before the final response is returned.
func WithStreaming(callback ModelStreamCallback) ExecutionOption {
return &executionOptions{Stream: callback}
}
// documentOptions are options for providing context documents to a prompt or generate request or as input to an embedder.
type documentOptions struct {
Documents []*Document // Docs to pass as context or input.
}
// DocumentOption is an option for providing context or input documents.
// It applies only to [Generate] and [prompt.Execute].
type DocumentOption interface {
applyDocument(*documentOptions) error
applyGenerate(*generateOptions) error
applyPromptExecute(*promptExecutionOptions) error
applyEmbedder(*embedderOptions) error
applyRetriever(*retrieverOptions) error
}
// applyDocument applies the option to the context options.
func (o *documentOptions) applyDocument(docOpts *documentOptions) error {
if o.Documents != nil {
if docOpts.Documents != nil {
return errors.New("cannot set documents more than once (WithDocs)")
}
docOpts.Documents = o.Documents
}
return nil
}
// applyGenerate applies the option to the generate options.
func (o *documentOptions) applyGenerate(genOpts *generateOptions) error {
return o.applyDocument(&genOpts.documentOptions)
}
// applyPromptExecute applies the option to the prompt generate options.
func (o *documentOptions) applyPromptExecute(pgOpts *promptExecutionOptions) error {
return o.applyDocument(&pgOpts.documentOptions)
}
// applyEmbedder applies the option to the embed options.
func (o *documentOptions) applyEmbedder(embedOpts *embedderOptions) error {
return o.applyDocument(&embedOpts.documentOptions)
}
// applyRetriever applies the option to the retrieve options.
func (o *documentOptions) applyRetriever(retOpts *retrieverOptions) error {
return o.applyDocument(&retOpts.documentOptions)
}
// WithTextDocs sets the text to be used as context documents for generation or as input to an embedder.
func WithTextDocs(text ...string) DocumentOption {
docs := make([]*Document, len(text))
for i, t := range text {
docs[i] = DocumentFromText(t, nil)
}
return &documentOptions{Documents: docs}
}
// WithDocs sets the documents to be used as context for generation or as input to an embedder.
func WithDocs(docs ...*Document) DocumentOption {
return &documentOptions{Documents: docs}
}
// evaluatorOptions are options for providing a dataset to evaluate.
type evaluatorOptions struct {
configOptions
Dataset []*Example // Dataset to evaluate.
ID string // ID of the evaluation.
Evaluator EvaluatorArg // Evaluator to use.
}
// EvaluatorOption is an option for providing a dataset to evaluate.
// It applies only to [Evaluator.Evaluate].
type EvaluatorOption interface {
applyEvaluator(*evaluatorOptions) error
}
// applyEvaluator applies the option to the evaluator options.
func (o *evaluatorOptions) applyEvaluator(evalOpts *evaluatorOptions) error {
if err := o.applyConfig(&evalOpts.configOptions); err != nil {
return err
}
if o.Dataset != nil {
if evalOpts.Dataset != nil {
return errors.New("cannot set dataset more than once (WithDataset)")
}
evalOpts.Dataset = o.Dataset
}
if o.ID != "" {
if evalOpts.ID != "" {
return errors.New("cannot set ID more than once (WithID)")
}
evalOpts.ID = o.ID
}
if o.Evaluator != nil {
if evalOpts.Evaluator != nil {
return errors.New("cannot set evaluator more than once (WithEvaluator or WithEvaluatorName)")
}
evalOpts.Evaluator = o.Evaluator
}
return nil
}
// WithDataset sets the dataset to do evaluation on.
func WithDataset(examples ...*Example) EvaluatorOption {
return &evaluatorOptions{Dataset: examples}
}
// WithID sets the ID of the evaluation to uniquely identify it.
func WithID(ID string) EvaluatorOption {
return &evaluatorOptions{ID: ID}
}
// WithEvaluator sets either a [Evaluator] or a [EvaluatorRef] that may contain a config.
// Passing [WithConfig] will take precedence over the config in WithEvaluator.
func WithEvaluator(evaluator EvaluatorArg) EvaluatorOption {
return &evaluatorOptions{Evaluator: evaluator}
}
// WithEvaluatorName sets the evaluator name to call for document evaluation.
// The evaluator name will be resolved to a [Evaluator] and may error if the reference is invalid.
func WithEvaluatorName(name string) EvaluatorOption {
return &evaluatorOptions{Evaluator: NewEvaluatorRef(name, nil)}
}
// embedderOptions holds configuration and input for an embedder request.
type embedderOptions struct {
configOptions
documentOptions
Embedder EmbedderArg // Embedder to use.
}
// EmbedderOption is an option for configuring an embedder request.
// It applies only to [Embed].
type EmbedderOption interface {
applyEmbedder(*embedderOptions) error
}
// applyEmbedder applies the option to the embed options.
func (o *embedderOptions) applyEmbedder(embedOpts *embedderOptions) error {
if err := o.applyConfig(&embedOpts.configOptions); err != nil {
return err
}
if err := o.applyDocument(&embedOpts.documentOptions); err != nil {
return err
}
if o.Embedder != nil {
if embedOpts.Embedder != nil {
return errors.New("cannot set embedder more than once (WithEmbedder or WithEmbedderName)")
}
embedOpts.Embedder = o.Embedder
}
return nil
}
// WithEmbedder sets either a [Embedder] or a [EmbedderRef] that may contain a config.
// Passing [WithConfig] will take precedence over the config in WithEmbedder.
func WithEmbedder(embedder EmbedderArg) EmbedderOption {
return &embedderOptions{Embedder: embedder}
}
// WithEmbedderName sets the embedder name to call for document embedding.
// The embedder name will be resolved to a [Embedder] and may error if the reference is invalid.
func WithEmbedderName(name string) EmbedderOption {
return &embedderOptions{Embedder: NewEmbedderRef(name, nil)}
}
// retrieverOptions holds configuration and input for a retriever request.
type retrieverOptions struct {
configOptions
documentOptions
Retriever RetrieverArg // Retriever to use.
}
// RetrieverOption is an option for configuring a retriever request.
// It applies only to [Retriever.Retrieve].
type RetrieverOption interface {
applyRetriever(*retrieverOptions) error
}
// applyRetriever applies the option to the retrieve options.
func (o *retrieverOptions) applyRetriever(retOpts *retrieverOptions) error {
if err := o.applyConfig(&retOpts.configOptions); err != nil {
return err
}
if err := o.applyDocument(&retOpts.documentOptions); err != nil {
return err
}
if o.Retriever != nil {
if retOpts.Retriever != nil {
return errors.New("cannot set retriever more than once (WithRetriever or WithRetrieverName)")
}
retOpts.Retriever = o.Retriever
}
return nil
}
// WithRetriever sets either a [Retriever] or a [RetrieverRef] that may contain a config.
// Passing [WithConfig] will take precedence over the config in WithRetriever.
func WithRetriever(retriever RetrieverArg) RetrieverOption {
return &retrieverOptions{Retriever: retriever}
}
// WithRetrieverName sets the retriever name to call for document retrieval.
// The retriever name will be resolved to a [Retriever] and may error if the reference is invalid.
func WithRetrieverName(name string) RetrieverOption {
return &retrieverOptions{Retriever: NewRetrieverRef(name, nil)}
}
// generateOptions are options for generating a model response by calling a model directly.
type generateOptions struct {
commonGenOptions
promptingOptions
outputOptions
executionOptions
documentOptions
RespondParts []*Part // Tool responses to return from interrupted tool calls.
RestartParts []*Part // Tool requests to restart interrupted tools with.
}
// GenerateOption is an option for generating a model response. It applies only to Generate().
type GenerateOption interface {
applyGenerate(*generateOptions) error
}
// applyGenerate applies the option to the generate options.
func (o *generateOptions) applyGenerate(genOpts *generateOptions) error {
if err := o.commonGenOptions.applyGenerate(genOpts); err != nil {
return err
}
if err := o.promptingOptions.applyGenerate(genOpts); err != nil {
return err
}
if err := o.outputOptions.applyGenerate(genOpts); err != nil {
return err
}
if err := o.executionOptions.applyGenerate(genOpts); err != nil {
return err
}
if err := o.documentOptions.applyGenerate(genOpts); err != nil {
return err
}
if o.RespondParts != nil {
if genOpts.RespondParts != nil {
return errors.New("cannot set respond parts more than once (WithToolResponses)")
}
genOpts.RespondParts = o.RespondParts
}
if o.RestartParts != nil {
if genOpts.RestartParts != nil {
return errors.New("cannot set restart parts more than once (WithToolRestarts)")
}
genOpts.RestartParts = o.RestartParts
}
return nil
}
// WithToolResponses sets the tool responses to return from interrupted tool calls.
func WithToolResponses(parts ...*Part) GenerateOption {
return &generateOptions{RespondParts: parts}
}
// WithToolRestarts sets the tool requests to restart interrupted tools with.
func WithToolRestarts(parts ...*Part) GenerateOption {
return &generateOptions{RestartParts: parts}
}
// WithResources specifies resources to be temporarily available during generation.
// Resources are unregistered resources that get attached to a temporary registry
// during the generation request and cleaned up afterward.
func WithResources(resources ...Resource) CommonGenOption {
return &withResources{resources: resources}
}
type withResources struct {
resources []Resource
}
func (w *withResources) applyCommonGen(o *commonGenOptions) error {
o.Resources = w.resources
return nil
}
func (w *withResources) applyPrompt(o *promptOptions) error {
return w.applyCommonGen(&o.commonGenOptions)
}
func (w *withResources) applyGenerate(o *generateOptions) error {
return w.applyCommonGen(&o.commonGenOptions)
}
func (w *withResources) applyPromptExecute(o *promptExecutionOptions) error {
return w.applyCommonGen(&o.commonGenOptions)
}
// promptExecutionOptions are options for generating a model response by executing a prompt.
type promptExecutionOptions struct {
commonGenOptions
executionOptions
documentOptions
Input any // Input fields for the prompt. If not nil this should be a struct that matches the prompt's input schema.
}
// PromptExecuteOption is an option for executing a prompt. It applies only to [prompt.Execute].
type PromptExecuteOption interface {
applyPromptExecute(*promptExecutionOptions) error
}
// applyPromptExecute applies the option to the prompt request options.
func (o *promptExecutionOptions) applyPromptExecute(pgOpts *promptExecutionOptions) error {
if err := o.commonGenOptions.applyPromptExecute(pgOpts); err != nil {
return err
}
if err := o.executionOptions.applyPromptExecute(pgOpts); err != nil {
return err
}
if err := o.documentOptions.applyPromptExecute(pgOpts); err != nil {
return err
}
if o.Input != nil {
if pgOpts.Input != nil {
return errors.New("cannot set input more than once (WithInput)")
}
pgOpts.Input = o.Input
}
return nil
}
// WithInput sets the input for the prompt request. Input must conform to the
// prompt's input schema and can either be a map[string]any or a struct of the same api.
func WithInput(input any) PromptExecuteOption {
return &promptExecutionOptions{Input: input}
}