Merge branch 'main' into pj/js-googleai-imagen-veo

Author: pavelgj

genkit-tools/common/src/types/eval.ts

@@ -23,7 +23,7 @@ import type {
   ListEvalKeysResponse,
   UpdateDatasetRequest,
 } from './apis';
-import { GenerateRequestSchema } from './model';
+import { GenerateActionOptionsSchema, GenerateRequestSchema } from './model';
 
 /**
  * This file defines schema and types that are used by the Eval store.
@@ -56,6 +56,17 @@ export const GenerateRequestJSONSchema = zodToJsonSchema(
   }
 ) as JSONSchema7;
 
+/**
+ * Combined GenerateInput JSON schema to support eval-inference using models
+ */
+export const GenerateInputJSONSchema = zodToJsonSchema(
+  z.union([GenerateRequestSchema, GenerateActionOptionsSchema]),
+  {
+    $refStrategy: 'none',
+    removeAdditionalStrategy: 'strict',
+  }
+) as JSONSchema7;
+
 /**
  * A single sample to be used for inference.
  **/

go/ai/generate.go

@@ -372,9 +372,29 @@ func Generate(ctx context.Context, r *registry.Registry, opts ...GenerateOption)
 		modelName = genOpts.ModelName
 	}
 
+	var dynamicTools []Tool
 	tools := make([]string, len(genOpts.Tools))
-	for i, tool := range genOpts.Tools {
-		tools[i] = tool.Name()
+	toolNames := make(map[string]bool)
+	for i, toolRef := range genOpts.Tools {
+		name := toolRef.Name()
+		// Redundant duplicate tool check with GenerateWithRequest otherwise we will panic when we register the dynamic tools.
+		if toolNames[name] {
+			return nil, core.NewError(core.INVALID_ARGUMENT, "ai.Generate: duplicate tool %q", name)
+		}
+		toolNames[name] = true
+		tools[i] = name
+		// Dynamic tools wouldn't have been registered by this point.
+		if LookupTool(r, name) == nil {
+			if tool, ok := toolRef.(Tool); ok {
+				dynamicTools = append(dynamicTools, tool)
+			}
+		}
+	}
+	if len(dynamicTools) > 0 {
+		r = r.NewChild()
+		for _, tool := range dynamicTools {
+			tool.Register(r)
+		}
 	}
 
 	messages := []*Message{}
@@ -596,7 +616,7 @@ func handleToolRequests(ctx context.Context, r *registry.Registry, req *ModelReq
 
 			output, err := tool.RunRaw(ctx, toolReq.Input)
 			if err != nil {
-				var tie *ToolInterruptError
+				var tie *toolInterruptError
 				if errors.As(err, &tie) {
 					logger.FromContext(ctx).Debug("tool %q triggered an interrupt: %v", toolReq.Name, tie.Metadata)
 
@@ -636,7 +656,7 @@ func handleToolRequests(ctx context.Context, r *registry.Registry, req *ModelReq
 	for range toolCount {
 		res := <-resultChan
 		if res.err != nil {
-			var tie *ToolInterruptError
+			var tie *toolInterruptError
 			if errors.As(res.err, &tie) {
 				hasInterrupts = true
 				continue
@@ -878,7 +898,7 @@ func handleResumedToolRequest(ctx context.Context, r *registry.Registry, genOpts
 
 				output, err := tool.RunRaw(resumedCtx, restartPart.ToolRequest.Input)
 				if err != nil {
-					var tie *ToolInterruptError
+					var tie *toolInterruptError
 					if errors.As(err, &tie) {
 						logger.FromContext(ctx).Debug("tool %q triggered an interrupt: %v", restartPart.ToolRequest.Name, tie.Metadata)
 

go/ai/generate_test.go

@@ -621,6 +621,124 @@ func TestGenerate(t *testing.T) {
 			t.Errorf("got text %q, want %q", res.Text(), expectedText)
 		}
 	})
+
+	t.Run("registers dynamic tools", func(t *testing.T) {
+		// Create a tool that is NOT registered in the global registry
+		dynamicTool := NewTool("dynamicTestTool", "a tool that is dynamically registered",
+			func(ctx *ToolContext, input struct {
+				Message string
+			}) (string, error) {
+				return "Dynamic: " + input.Message, nil
+			},
+		)
+
+		// Verify the tool is not in the global registry
+		if LookupTool(r, "dynamicTestTool") != nil {
+			t.Fatal("dynamicTestTool should not be registered in global registry")
+		}
+
+		// Create a model that will call the dynamic tool then provide a final response
+		roundCount := 0
+		info := &ModelInfo{
+			Supports: &ModelSupports{
+				Multiturn: true,
+				Tools:     true,
+			},
+		}
+		toolCallModel := DefineModel(r, "test", "toolcall", info,
+			func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) {
+				roundCount++
+				if roundCount == 1 {
+					// First response: call the dynamic tool
+					return &ModelResponse{
+						Request: gr,
+						Message: &Message{
+							Role: RoleModel,
+							Content: []*Part{
+								NewToolRequestPart(&ToolRequest{
+									Name:  "dynamicTestTool",
+									Input: map[string]any{"Message": "Hello from dynamic tool"},
+								}),
+							},
+						},
+					}, nil
+				}
+				// Second response: provide final answer based on tool response
+				var toolResult string
+				for _, msg := range gr.Messages {
+					if msg.Role == RoleTool {
+						for _, part := range msg.Content {
+							if part.ToolResponse != nil {
+								toolResult = part.ToolResponse.Output.(string)
+							}
+						}
+					}
+				}
+				return &ModelResponse{
+					Request: gr,
+					Message: &Message{
+						Role: RoleModel,
+						Content: []*Part{
+							NewTextPart(toolResult),
+						},
+					},
+				}, nil
+			})
+
+		// Use Generate with the dynamic tool - this should trigger the dynamic registration
+		res, err := Generate(context.Background(), r,
+			WithModel(toolCallModel),
+			WithPrompt("call the dynamic tool"),
+			WithTools(dynamicTool),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// The tool should have been called and returned a response
+		expectedText := "Dynamic: Hello from dynamic tool"
+		if res.Text() != expectedText {
+			t.Errorf("expected text %q, got %q", expectedText, res.Text())
+		}
+
+		// Verify two rounds were executed: tool call + final response
+		if roundCount != 2 {
+			t.Errorf("expected 2 rounds, got %d", roundCount)
+		}
+
+		// Verify the tool is still not in the global registry (it was registered in a child)
+		if LookupTool(r, "dynamicTestTool") != nil {
+			t.Error("dynamicTestTool should not be registered in global registry after generation")
+		}
+	})
+
+	t.Run("handles duplicate dynamic tools", func(t *testing.T) {
+		// Create two tools with the same name
+		dynamicTool1 := NewTool("duplicateTool", "first tool",
+			func(ctx *ToolContext, input any) (string, error) {
+				return "tool1", nil
+			},
+		)
+		dynamicTool2 := NewTool("duplicateTool", "second tool",
+			func(ctx *ToolContext, input any) (string, error) {
+				return "tool2", nil
+			},
+		)
+
+		// Using both tools should result in an error
+		_, err := Generate(context.Background(), r,
+			WithModel(echoModel),
+			WithPrompt("test duplicate tools"),
+			WithTools(dynamicTool1, dynamicTool2),
+		)
+
+		if err == nil {
+			t.Fatal("expected error for duplicate tool names")
+		}
+		if !strings.Contains(err.Error(), "duplicate tool \"duplicateTool\"") {
+			t.Errorf("unexpected error message: %v", err)
+		}
+	})
 }
 
 func TestModelVersion(t *testing.T) {

go/ai/option.go

@@ -47,7 +47,6 @@ type ConfigOption interface {
 	applyEmbedder(*embedderOptions) error
 	applyRetriever(*retrieverOptions) error
 	applyEvaluator(*evaluatorOptions) error
-	applyIndexer(*indexerOptions) error
 }
 
 // applyConfig applies the option to the config options.
@@ -96,11 +95,6 @@ func (o *configOptions) applyEvaluator(opts *evaluatorOptions) error {
 	return o.applyConfig(&opts.configOptions)
 }
 
-// applyIndexer applies the option to the indexer options.
-func (o *configOptions) applyIndexer(opts *indexerOptions) error {
-	return o.applyConfig(&opts.configOptions)
-}
-
 // WithConfig sets the configuration.
 func WithConfig(config any) ConfigOption {
 	return &configOptions{Config: config}
@@ -580,7 +574,6 @@ type DocumentOption interface {
 	applyPromptExecute(*promptExecutionOptions) error
 	applyEmbedder(*embedderOptions) error
 	applyRetriever(*retrieverOptions) error
-	applyIndexer(*indexerOptions) error
 }
 
 // applyDocument applies the option to the context options.
@@ -615,11 +608,6 @@ func (o *documentOptions) applyRetriever(retOpts *retrieverOptions) error {
 	return o.applyDocument(&retOpts.documentOptions)
 }
 
-// applyIndexer applies the option to the indexer options.
-func (o *documentOptions) applyIndexer(idxOpts *indexerOptions) error {
-	return o.applyDocument(&idxOpts.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))
@@ -730,31 +718,6 @@ func (o *retrieverOptions) applyRetriever(retOpts *retrieverOptions) error {
 	return nil
 }
 
-// indexerOptions holds configuration and input for an embedder request.
-type indexerOptions struct {
-	configOptions
-	documentOptions
-}
-
-// IndexerOption is an option for configuring an embedder request.
-// It applies only to [Index].
-type IndexerOption interface {
-	applyIndexer(*indexerOptions) error
-}
-
-// applyIndexer applies the option to the indexer options.
-func (o *indexerOptions) applyIndexer(idxOpts *indexerOptions) error {
-	if err := o.applyConfig(&idxOpts.configOptions); err != nil {
-		return err
-	}
-
-	if err := o.applyDocument(&idxOpts.documentOptions); err != nil {
-		return err
-	}
-
-	return nil
-}
-
 // generateOptions are options for generating a model response by calling a model directly.
 type generateOptions struct {
 	commonGenOptions