diff --git a/agent/agent.go b/agent/agent.go index 3b194be07..d44304319 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -48,6 +48,12 @@ type Agent interface { // New creates an Agent with a custom logic defined by Run function. func New(cfg Config) (Agent, error) { + if cfg.ValidateFunc != nil { + if err := cfg.ValidateFunc(); err != nil { + return nil, fmt.Errorf("agent %q validation failed: %w", cfg.Name, err) + } + } + subAgentSet := make(map[Agent]bool) for _, subAgent := range cfg.SubAgents { if _, ok := subAgentSet[subAgent]; ok { @@ -99,6 +105,10 @@ type Config struct { // created from the content or error of that callback and the remaining // callbacks will be skipped. AfterAgentCallbacks []AfterAgentCallback + + // ValidateFunc is an optional function that checks if the agent is + // configured correctly. + ValidateFunc func() error } // Artifacts interface provides methods to work with artifacts of the current diff --git a/agent/llmagent/llmagent.go b/agent/llmagent/llmagent.go index b1c7254fe..d716b40ef 100644 --- a/agent/llmagent/llmagent.go +++ b/agent/llmagent/llmagent.go @@ -32,6 +32,12 @@ import ( // New is a constructor for LLMAgent. func New(cfg Config) (agent.Agent, error) { + if cfg.ValidateFunc != nil { + if err := cfg.ValidateFunc(); err != nil { + return nil, fmt.Errorf("llmagent %q validation failed: %w", cfg.Name, err) + } + } + beforeModelCallbacks := make([]llminternal.BeforeModelCallback, 0, len(cfg.BeforeModelCallbacks)) for _, c := range cfg.BeforeModelCallbacks { beforeModelCallbacks = append(beforeModelCallbacks, llminternal.BeforeModelCallback(c)) @@ -255,6 +261,10 @@ type Config struct { // - Extracts agent reply for later use, such as in tools, callbacks, etc. // - Connects agents to coordinate with each other. OutputKey string + + // ValidateFunc is an optional function that checks if the agent is + // configured correctly. + ValidateFunc func() error } // BeforeModelCallback that is called before sending a request to the model. diff --git a/internal/toolinternal/context.go b/internal/toolinternal/context.go index 9a3525d84..3d4b62f98 100644 --- a/internal/toolinternal/context.go +++ b/internal/toolinternal/context.go @@ -16,6 +16,7 @@ package toolinternal import ( "context" + "errors" "github.com/google/uuid" "google.golang.org/genai" @@ -28,12 +29,18 @@ import ( "google.golang.org/adk/tool" ) +// ErrArtifactServiceNotConfigured is returned when artifact service operations are attempted without configuration. +var ErrArtifactServiceNotConfigured = errors.New("artifact service not configured") + type internalArtifacts struct { agent.Artifacts eventActions *session.EventActions } func (ia *internalArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + if ia == nil { + return nil, ErrArtifactServiceNotConfigured + } resp, err := ia.Artifacts.Save(ctx, name, data) if err != nil { return resp, err @@ -48,6 +55,20 @@ func (ia *internalArtifacts) Save(ctx context.Context, name string, data *genai. return resp, nil } +func (ia *internalArtifacts) List(ctx context.Context) (*artifact.ListResponse, error) { + if ia == nil { + return nil, ErrArtifactServiceNotConfigured + } + return ia.Artifacts.List(ctx) +} + +func (ia *internalArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { + if ia == nil { + return nil, ErrArtifactServiceNotConfigured + } + return ia.Artifacts.Load(ctx, name) +} + func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions *session.EventActions) tool.Context { if functionCallID == "" { functionCallID = uuid.NewString() @@ -60,15 +81,21 @@ func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions } cbCtx := contextinternal.NewCallbackContextWithDelta(ctx, actions.StateDelta) + // Only create internalArtifacts if the underlying Artifacts service is configured + var artifacts *internalArtifacts + if ctx.Artifacts() != nil { + artifacts = &internalArtifacts{ + Artifacts: ctx.Artifacts(), + eventActions: actions, + } + } + return &toolContext{ CallbackContext: cbCtx, invocationContext: ctx, functionCallID: functionCallID, eventActions: actions, - artifacts: &internalArtifacts{ - Artifacts: ctx.Artifacts(), - eventActions: actions, - }, + artifacts: artifacts, } } diff --git a/internal/toolinternal/context_test.go b/internal/toolinternal/context_test.go index 517cf6260..7ed83274c 100644 --- a/internal/toolinternal/context_test.go +++ b/internal/toolinternal/context_test.go @@ -15,6 +15,7 @@ package toolinternal import ( + "errors" "testing" "google.golang.org/adk/agent" @@ -36,3 +37,45 @@ func TestToolContext(t *testing.T) { t.Errorf("ToolContext(%+T) is unexpectedly an InvocationContext", got) } } + +func TestInternalArtifacts_NilSafe(t *testing.T) { + // Create invocation context without artifact service + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{ + Artifacts: nil, + }) + toolCtx := NewToolContext(inv, "fn1", &session.EventActions{}) + + artifacts := toolCtx.Artifacts() + // artifacts will be nil when service not configured + + tests := []struct { + name string + call func() (any, error) + }{ + { + name: "List", + call: func() (any, error) { return artifacts.List(t.Context()) }, + }, + { + name: "Load", + call: func() (any, error) { return artifacts.Load(t.Context(), "test.txt") }, + }, + { + name: "Save", + call: func() (any, error) { return artifacts.Save(t.Context(), "test.txt", nil) }, + }, + } + + for _, tt := range tests { + t.Run(tt.name+" returns error", func(t *testing.T) { + _, err := tt.call() + if err == nil { + t.Error("Expected an error, got nil") + return + } + if !errors.Is(err, ErrArtifactServiceNotConfigured) { + t.Errorf("Expected ErrArtifactServiceNotConfigured, got: %v", err) + } + }) + } +} diff --git a/runner/runner.go b/runner/runner.go index db56cd5ad..412e844c5 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -65,6 +65,11 @@ func New(cfg Config) (*Runner, error) { return nil, fmt.Errorf("failed to create agent tree: %w", err) } + // Validate that required services are configured for tools + if err := validateConfiguration(cfg); err != nil { + return nil, err + } + return &Runner{ appName: cfg.AppName, rootAgent: cfg.Agent, @@ -268,3 +273,48 @@ func findAgent(curAgent agent.Agent, targetName string) agent.Agent { } return nil } + +// Validator is an optional interface that Agents and Tools can implement +// to validate if the runner configuration meets their requirements. +type Validator interface { + Validate(Config) error +} + +// validateConfiguration checks that required services are available for tools. +func validateConfiguration(cfg Config) error { + return walkAgentTree(cfg.Agent, func(a agent.Agent) error { + if v, ok := a.(Validator); ok { + if err := v.Validate(cfg); err != nil { + return fmt.Errorf("agent %q validation failed: %w", a.Name(), err) + } + } + + llmAgent, ok := a.(llminternal.Agent) + if !ok { + return nil + } + + state := llminternal.Reveal(llmAgent) + for _, t := range state.Tools { + if v, ok := t.(Validator); ok { + if err := v.Validate(cfg); err != nil { + return fmt.Errorf("agent %q tool %q validation failed: %w", a.Name(), t.Name(), err) + } + } + } + return nil + }) +} + +// walkAgentTree recursively walks the agent tree and applies fn to each agent. +func walkAgentTree(a agent.Agent, fn func(agent.Agent) error) error { + if err := fn(a); err != nil { + return err + } + for _, sub := range a.SubAgents() { + if err := walkAgentTree(sub, fn); err != nil { + return err + } + } + return nil +} diff --git a/runner/runner_test.go b/runner/runner_test.go index 2249edebe..1f3710443 100644 --- a/runner/runner_test.go +++ b/runner/runner_test.go @@ -22,12 +22,12 @@ import ( "strings" "testing" - "google.golang.org/genai" - "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" "google.golang.org/adk/artifact" "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/genai" ) func TestRunner_findAgentToRun(t *testing.T) { @@ -314,6 +314,172 @@ func TestRunner_SaveInputBlobsAsArtifacts(t *testing.T) { } } +func TestNew_ValidatesAgent(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + validateErr error + wantErr bool + errContains string + }{ + { + name: "agent validation passes", + validateErr: nil, + wantErr: false, + }, + { + name: "agent validation fails", + validateErr: fmt.Errorf("validation failed"), + wantErr: true, + errContains: "llmagent \"validating_agent\" validation failed: validation failed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := llmagent.New(llmagent.Config{ + Name: "validating_agent", + ValidateFunc: func() error { + return tt.validateErr + }, + }) + + if tt.wantErr { + if err == nil { + t.Errorf("llmagent.New() expected error but got nil") + return + } + if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) { + t.Errorf("llmagent.New() error = %v, want error containing %q", err, tt.errContains) + } + } else { + if err != nil { + t.Errorf("llmagent.New() unexpected error = %v", err) + } + } + }) + } +} + +type mockValidatorTool struct { + name string + validateFunc func(Config) error +} + +func (t *mockValidatorTool) Name() string { return t.name } +func (t *mockValidatorTool) Description() string { return "mock tool" } +func (t *mockValidatorTool) IsLongRunning() bool { return false } +func (t *mockValidatorTool) Run(ctx tool.Context, args any) (map[string]any, error) { return nil, nil } +func (t *mockValidatorTool) Validate(cfg Config) error { + if t.validateFunc != nil { + return t.validateFunc(cfg) + } + return nil +} + +func TestNew_ValidatesTool(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + agent agent.Agent + artifactService artifact.Service + wantErr bool + errContains string + }{ + { + name: "error when tool validation fails", + agent: must(llmagent.New(llmagent.Config{ + Name: "test_agent", + Tools: []tool.Tool{ + &mockValidatorTool{ + name: "validation_tool", + validateFunc: func(cfg Config) error { + if cfg.ArtifactService == nil { + return fmt.Errorf("requires ArtifactService") + } + return nil + }, + }, + }, + })), + artifactService: nil, + wantErr: true, + errContains: "requires ArtifactService", + }, + { + name: "ok when tool validation passes", + agent: must(llmagent.New(llmagent.Config{ + Name: "test_agent", + Tools: []tool.Tool{ + &mockValidatorTool{ + name: "validation_tool", + validateFunc: func(cfg Config) error { + if cfg.ArtifactService == nil { + return fmt.Errorf("requires ArtifactService") + } + return nil + }, + }, + }, + })), + artifactService: artifact.InMemoryService(), + wantErr: false, + }, + { + name: "error when nested tool validation fails", + agent: must(llmagent.New(llmagent.Config{ + Name: "parent_agent", + SubAgents: []agent.Agent{ + must(llmagent.New(llmagent.Config{ + Name: "child_agent", + Tools: []tool.Tool{ + &mockValidatorTool{ + name: "validation_tool", + validateFunc: func(cfg Config) error { + if cfg.ArtifactService == nil { + return fmt.Errorf("requires ArtifactService") + } + return nil + }, + }, + }, + })), + }, + })), + artifactService: nil, + wantErr: true, + errContains: "child_agent", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := New(Config{ + AppName: "testApp", + Agent: tt.agent, + SessionService: session.InMemoryService(), + ArtifactService: tt.artifactService, + }) + + if tt.wantErr { + if err == nil { + t.Errorf("New() expected error but got nil") + return + } + if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) { + t.Errorf("New() error = %v, want error containing %q", err, tt.errContains) + } + } else { + if err != nil { + t.Errorf("New() unexpected error = %v", err) + } + } + }) + } +} + // creates agentTree for tests and returns references to the agents func agentTree(t *testing.T) agentTreeStruct { t.Helper() diff --git a/tool/loadartifactstool/load_artifacts_tool.go b/tool/loadartifactstool/load_artifacts_tool.go index 80b034777..29c5fcf6c 100644 --- a/tool/loadartifactstool/load_artifacts_tool.go +++ b/tool/loadartifactstool/load_artifacts_tool.go @@ -29,6 +29,7 @@ import ( "google.golang.org/adk/internal/toolinternal/toolutils" "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" + "google.golang.org/adk/runner" "google.golang.org/adk/tool" ) @@ -61,6 +62,14 @@ func (t *artifactsTool) IsLongRunning() bool { return false } +// Validate implements runner.Validator. +func (t *artifactsTool) Validate(cfg runner.Config) error { + if cfg.ArtifactService == nil { + return fmt.Errorf("tool %q requires ArtifactService to be configured in runner", t.name) + } + return nil +} + // Declaration returns the GenAI FunctionDeclaration for the load_artifacts tool. // // This declaration allows the LLM to understand and call the tool diff --git a/tool/loadartifactstool/load_artifacts_tool_test.go b/tool/loadartifactstool/load_artifacts_tool_test.go index fd72d0066..92f31d085 100644 --- a/tool/loadartifactstool/load_artifacts_tool_test.go +++ b/tool/loadartifactstool/load_artifacts_tool_test.go @@ -15,6 +15,7 @@ package loadartifactstool_test import ( + "errors" "strings" "testing" @@ -277,6 +278,32 @@ func TestLoadArtifactsTool_ProcessRequest_Artifacts_OtherFunctionCall(t *testing } } +func TestLoadArtifactsTool_ProcessRequest_NoArtifactService(t *testing.T) { + loadArtifactsTool := loadartifactstool.New() + + // Create tool context WITHOUT artifact service configured + ctx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ + Artifacts: nil, // No artifact service + }) + tc := toolinternal.NewToolContext(ctx, "", nil) + + llmRequest := &model.LLMRequest{} + + requestProcessor, ok := loadArtifactsTool.(toolinternal.RequestProcessor) + if !ok { + t.Fatal("loadArtifactsTool does not implement RequestProcessor") + } + + err := requestProcessor.ProcessRequest(tc, llmRequest) + if err == nil { + t.Fatal("Expected error when artifact service not configured, got nil") + } + + if !errors.Is(err, toolinternal.ErrArtifactServiceNotConfigured) { + t.Errorf("Expected ErrArtifactServiceNotConfigured, got: %v", err) + } +} + func createToolContext(t *testing.T) tool.Context { t.Helper()