ronantakizawa/codereview-bench huggingface dataset:
Repository: google/adk-go
File: agent/workflowagents/parallelagent/agent_test.go (Go)

Code before:
			}

			if tt.wantEvents != nil {
				eventCompareFunc := func(e1, e2 *session.Event) int {
					if e1.Author <= e2.Author {
						return -1
					}
					if e1.Author == e2.Author {
						return 0
					}
					return 1
				}

				slices.SortFunc(tt.wantEvents, eventCompareFunc)
				slices.SortFunc(gotEvents, eventCompareFunc)

				if diff := cmp.Diff(tt.wantEvents, gotEvents); diff != "" {
					t.Errorf("events mismatch (-want +got):\n%s", diff)
				}
			}
		})
	}
}

// newParallelAgent creates parallel agent with 2 subagents emitting maxIterations events or infinitely if maxIterations==0.
func newParallelAgent(t *testing.T, maxIterations uint, agentErr error) agent.Agent {
	subAgents := []agent.Agent{
		must(loopagent.New(loopagent.Config{
			MaxIterations: maxIterations,
			AgentConfig: agent.Config{
				Name: "loop1",
				SubAgents: []agent.Agent{
					must(agent.New(agent.Config{
						Name: "sub1",
						Run:  customRun(1, agentErr),
					},
					)),
				},
			},
		})),
		must(loopagent.New(loopagent.Config{
			MaxIterations: maxIterations,
			AgentConfig: agent.Config{
				Name: "loop2",
				SubAgents: []agent.Agent{
					must(agent.New(agent.Config{
						Name: "sub2",
						Run:  customRun(2, agentErr),
					},
					)),
				},

Review comment (suggestion):
is it possible to have a test case where one of the agents terminates early with error while others keep running (to test the other agents stops). And I wonder if # agents >> 2, that will increase the chance of race the race detector can catch. 

Code after:
					}
					if tt.agentError != nil && err == nil {
						// In case of agent error some events from other agents can be processed before error is returned.
						continue
					}
					t.Errorf("got unexpected error: %v", err)
				}

				gotEvents = append(gotEvents, event)
			}

			if tt.wantEvents != nil {
				eventCompareFunc := func(e1, e2 *session.Event) int {
					if e1.Author <= e2.Author {
						return -1
					}
					if e1.Author == e2.Author {
						return 0
					}
					return 1
				}

				slices.SortFunc(tt.wantEvents, eventCompareFunc)
				slices.SortFunc(gotEvents, eventCompareFunc)

				if diff := cmp.Diff(tt.wantEvents, gotEvents); diff != "" {
					t.Errorf("events mismatch (-want +got):\n%s", diff)
				}
			}
		})
	}
}

// newParallelAgent creates parallel agent with 2 subagents emitting maxIterations events or infinitely if maxIterations==0.
func newParallelAgent(t *testing.T, maxIterations uint, numSubAgents int, agentErr error) agent.Agent {
	var subAgents []agent.Agent

	for i := 1; i <= numSubAgents; i++ {
		subAgents = append(subAgents, must(loopagent.New(loopagent.Config{
			MaxIterations: maxIterations,
			AgentConfig: agent.Config{
				Name: fmt.Sprintf("loop_agent_%d", i),
				SubAgents: []agent.Agent{
					must(agent.New(agent.Config{
						Name: fmt.Sprintf("sub%d", i),
						Run:  customRun(i, nil),
					},
					)),
				},
			},
		})))
Repository: google/adk-go
File: memoryservice/inmemory.go (Go)

Code before:
	store map[key]map[sessionID][]value
}

func (s *inMemoryService) AddSession(ctx context.Context, curSession session.Session) error {
	var values []value

	for event := range curSession.Events().All() {
		if event.LLMResponse == nil || event.LLMResponse.Content == nil {
			continue
		}

		words := make(map[string]struct{})
		for _, part := range event.LLMResponse.Content.Parts {
			if part.Text == "" {
				continue
			}

			maps.Copy(words, extractWords(part.Text))
		}

		if len(words) == 0 {
			continue
		}

		values = append(values, value{
			event: event,
			words: words,
		})
	}

	k := key{
		appName: curSession.ID().AppName,
		userID:  curSession.ID().UserID,
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	v, ok := s.store[k]
	if !ok {
		s.store[k] = map[sessionID][]value{
			curSession.ID().SessionID: values,
		}
		return nil
	}

	v[curSession.ID().SessionID] = values
	return nil
}

func (s *inMemoryService) SearchMemory(ctx context.Context, req *SearchMemoryRequest) (*SearchMemoryResponse, error) {

Review comment (suggestion):
optional: Makes me wonder if we want to keep the entire event if the SearchMemory API needs only the LLMResponse.Content. 

Code after:
type inMemoryService struct {
	mu    sync.RWMutex
	store map[key]map[sessionID][]value
}

func (s *inMemoryService) AddSession(ctx context.Context, curSession session.Session) error {
	var values []value

	for event := range curSession.Events().All() {
		if event.LLMResponse == nil || event.LLMResponse.Content == nil {
			continue
		}

		words := make(map[string]struct{})
		for _, part := range event.LLMResponse.Content.Parts {
			if part.Text == "" {
				continue
			}

			maps.Copy(words, extractWords(part.Text))
		}

		if len(words) == 0 {
			continue
		}

		values = append(values, value{
			content:   event.LLMResponse.Content,
			author:    event.Author,
			timestamp: event.Time,
			words:     words,
		})
	}

	k := key{
		appName: curSession.ID().AppName,
		userID:  curSession.ID().UserID,
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	v, ok := s.store[k]
	if !ok {
		v = map[sessionID][]value{}
		s.store[k] = v
	}

	v[curSession.ID().SessionID] = values
	return nil
}
Repository: google/adk-go
File: memoryservice/inmemory.go (Go)

Code before:
		}
	}

	return res, nil
}

func checkMapsIntersect(m1, m2 map[string]struct{}) bool {
	if len(m1) == 0 || len(m2) == 0 {
		return false
	}

	// Iterate over the smaller map.
	if len(m1) > len(m2) {
		m1, m2 = m2, m1
	}

	for k := range m1 {
		if _, ok := m2[k]; ok {
			return true
		}
	}

	return false
}

var wordRegex = regexp.MustCompile(`[A-Za-z]+`)

func extractWords(text string) map[string]struct{} {
	res := make(map[string]struct{})

	for _, word := range wordRegex.FindAllString(text, -1) {
		res[strings.ToLower(word)] = struct{}{}
	}

	return res
}


Review comment (suggestion):
Does the in-memory memory service allows only english alphabet text search?
I wonder if we just split (strings.SplitSeq) based on space and index everything.  

Code after:
				})
			}
		}
	}

	return res, nil
}

func checkMapsIntersect(m1, m2 map[string]struct{}) bool {
	if len(m1) == 0 || len(m2) == 0 {
		return false
	}

	// Iterate over the smaller map.
	if len(m1) > len(m2) {
		m1, m2 = m2, m1
	}

	for k := range m1 {
		if _, ok := m2[k]; ok {
			return true
		}
	}

	return false
}

func extractWords(text string) map[string]struct{} {
	res := make(map[string]struct{})

	for s := range strings.SplitSeq(text, " ") {
		if s == "" {
			continue
		}
		res[strings.ToLower(s)] = struct{}{}
	}

	return res
}
Repository: google/adk-go
File: internal/artifactsinternal/artifacts.go (Go)

Code before:
//
// 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.

package artifactsinternal

import (
	"context"

	"google.golang.org/adk/agent"
	"google.golang.org/adk/artifactservice"
	"google.golang.org/adk/session"
	"google.golang.org/genai"
)

// artifacts implements Artifacts
type artifacts struct {
	service artifactservice.Service
	id      session.ID
}

// NewArtifacts creates and returns Artifacts implementation.
func NewArtifacts(service artifactservice.Service, id session.ID) *artifacts {
	return &artifacts{
		service: service,
		id:      id,
	}
}

func (a *artifacts) Save(name string, data genai.Part) error {
	_, err := a.service.Save(context.Background(), &artifactservice.SaveRequest{
		AppName:   a.id.AppName,
		UserID:    a.id.UserID,
		SessionID: a.id.SessionID,
		FileName:  name,
		Part:      &data,
	})
	return err
}

func (a *artifacts) Load(name string) (genai.Part, error) {
	loadResponse, err := a.service.Load(context.Background(), &artifactservice.LoadRequest{
		AppName:   a.id.AppName,
		UserID:    a.id.UserID,
		SessionID: a.id.SessionID,
		FileName:  name,
	})
	if err != nil {

Review comment (suggestion):
Since this package is now internal, we can name `artifacts` as `Artifacts` and make this function return param an exported type.


Code after:
//
// 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.

package artifactsinternal

import (
	"context"

	"google.golang.org/adk/agent"
	"google.golang.org/adk/artifactservice"
	"google.golang.org/adk/session"
	"google.golang.org/genai"
)

// Artifacts implements Artifacts
type Artifacts struct {
	service artifactservice.Service
	id      session.ID
}

// NewArtifacts creates and returns Artifacts implementation.
func NewArtifacts(service artifactservice.Service, id session.ID) *Artifacts {
	return &Artifacts{
		service: service,
		id:      id,
	}
}

func (a *Artifacts) Save(name string, data genai.Part) error {
	_, err := a.service.Save(context.Background(), &artifactservice.SaveRequest{
		AppName:   a.id.AppName,
		UserID:    a.id.UserID,
		SessionID: a.id.SessionID,
		FileName:  name,
		Part:      &data,
	})
	return err
}

func (a *Artifacts) Load(name string) (genai.Part, error) {
	loadResponse, err := a.service.Load(context.Background(), &artifactservice.LoadRequest{
		AppName:   a.id.AppName,
		UserID:    a.id.UserID,
		SessionID: a.id.SessionID,
		FileName:  name,
	})
	if err != nil {
Repository: google/adk-go
File: tool/load_artifacts_tool.go (Go)

Code before:
}

// Run implements tool.Tool.
func (t *loadArtifactsTool) Run(ctx Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("unexpected args type, got: %T", args)
	}

	var artifactNames []string
	artifactNamesRaw, exists := m["artifact_names"]
	if !exists {
		artifactNames = []string{}
	} else {
		switch v := artifactNamesRaw.(type) {
		case []string:
			artifactNames = v
		case []any:
			for _, item := range v {
				strItem, ok := item.(string)
				if !ok {
					return nil, fmt.Errorf("element in artifact_names slice is not a string, got type %T for value %v", item, item)
				} else {
					artifactNames = append(artifactNames, strItem)
				}
			}
		}
	}
	result := map[string]any{
		"artifact_names": artifactNames,
	}
	return result, nil
}

// ProcessRequest implements tool.Tool.
func (t *loadArtifactsTool) ProcessRequest(ctx Context, req *llm.Request) error {
	if req.Tools == nil {
		req.Tools = make(map[string]any)
	}
	name := t.Name()
	if _, ok := req.Tools[name]; ok {
		return fmt.Errorf("duplicate tool: %q", name)
	}
	req.Tools[name] = t

	if req.GenerateConfig == nil {
		req.GenerateConfig = &genai.GenerateContentConfig{}
	}
	if decl := t.Declaration(); decl != nil {
		req.GenerateConfig.Tools = append(req.GenerateConfig.Tools, &genai.Tool{
			FunctionDeclarations: []*genai.FunctionDeclaration{decl},

Review comment (suggestion):
Shouldn't we handle types other than []string/[]any?

Another approach is marshal to json&unmarshal to []string. 

Code after:
}

// Run implements tool.Tool.
func (t *loadArtifactsTool) Run(ctx Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("unexpected args type, got: %T", args)
	}
	var artifactNames []string
	artifactNamesRaw, exists := m["artifact_names"]
	if !exists {
		artifactNames = []string{}
	} else {
		// In order to cast properly from []any to []string we're gonna marshal and then
		// unmarshal the artifact_names value.
		artifactNamesJson, err := json.Marshal(artifactNamesRaw)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal artifact_names to JSON: %w", err)
		}
		if err := json.Unmarshal(artifactNamesJson, &artifactNames); err != nil {
			return nil, fmt.Errorf("failed to unmarshal artifact_names from JSON to []string: %w", err)
		}
		// Ensure the slice is not nil if it's empty
		if artifactNames == nil {
			artifactNames = []string{}
		}
	}
	result := map[string]any{
		"artifact_names": artifactNames,
	}
	return result, nil
}

// ProcessRequest implements tool.Tool.
func (t *loadArtifactsTool) ProcessRequest(ctx Context, req *llm.Request) error {
	if req.Tools == nil {
		req.Tools = make(map[string]any)
	}
	name := t.Name()
	if _, ok := req.Tools[name]; ok {
		return fmt.Errorf("duplicate tool: %q", name)
	}
	req.Tools[name] = t

	if req.GenerateConfig == nil {
		req.GenerateConfig = &genai.GenerateContentConfig{}
	}
	if decl := t.Declaration(); decl != nil {
		req.GenerateConfig.Tools = append(req.GenerateConfig.Tools, &genai.Tool{
			FunctionDeclarations: []*genai.FunctionDeclaration{decl},
		})
Repository: google/adk-go
File: tool/load_artifacts_tool.go (Go)

Code before:
		Parameters: &genai.Schema{
			Type: "OBJECT",
			Properties: map[string]*genai.Schema{
				"artifact_names": {
					Type: "ARRAY",
					Items: &genai.Schema{
						Type: "STRING",
					},
				},
			},
		},
	}
}

// Run implements tool.Tool.
func (t *loadArtifactsTool) Run(ctx Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("unexpected args type, got: %T", args)
	}
	var artifactNames []string
	artifactNamesRaw, exists := m["artifact_names"]
	if !exists {
		artifactNames = []string{}
	} else {
		// In order to cast porperly from []any to []string we're gonna marshal and then
		// unmarshal the artifact_names value.
		artifactNamesJson, err := json.Marshal(artifactNamesRaw)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal artifact_names to JSON: %w", err)
		}
		if err := json.Unmarshal(artifactNamesJson, &artifactNames); err != nil {
			return nil, fmt.Errorf("failed to unmarshal artifact_names from JSON to []string: %w", err)
		}
		// Ensure the slice is not nil if it's empty
		if artifactNames == nil {
			artifactNames = []string{}
		}
	}
	result := map[string]any{
		"artifact_names": artifactNames,
	}
	return result, nil
}

// ProcessRequest implements tool.Tool.
func (t *loadArtifactsTool) ProcessRequest(ctx Context, req *llm.Request) error {
	if req.Tools == nil {
		req.Tools = make(map[string]any)
	}
	name := t.Name()

Review comment (suggestion):
```suggestion
		// In order to cast properly from []any to []string we're gonna marshal and then
```

Code after:
		Parameters: &genai.Schema{
			Type: "OBJECT",
			Properties: map[string]*genai.Schema{
				"artifact_names": {
					Type: "ARRAY",
					Items: &genai.Schema{
						Type: "STRING",
					},
				},
			},
		},
	}
}

// Run implements tool.Tool.
func (t *loadArtifactsTool) Run(ctx Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("unexpected args type, got: %T", args)
	}
	var artifactNames []string
	artifactNamesRaw, exists := m["artifact_names"]
	if !exists {
		artifactNames = []string{}
	} else {
		// In order to cast properly from []any to []string we're gonna marshal and then
		// unmarshal the artifact_names value.
		artifactNamesJson, err := json.Marshal(artifactNamesRaw)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal artifact_names to JSON: %w", err)
		}
		if err := json.Unmarshal(artifactNamesJson, &artifactNames); err != nil {
			return nil, fmt.Errorf("failed to unmarshal artifact_names from JSON to []string: %w", err)
		}
		// Ensure the slice is not nil if it's empty
		if artifactNames == nil {
			artifactNames = []string{}
		}
	}
	result := map[string]any{
		"artifact_names": artifactNames,
	}
	return result, nil
}

// ProcessRequest implements tool.Tool.
func (t *loadArtifactsTool) ProcessRequest(ctx Context, req *llm.Request) error {
	if req.Tools == nil {
		req.Tools = make(map[string]any)
	}
	name := t.Name()
Repository: google/adk-go
File: cmd/adkgo/root/build/webui/webui.go (Go)

Code before:

func (f *runLocalFlags) cleanTemp() error {
	err := util.LogStartStop("Cleaning target directory",
		func(p util.Printer) error {
			p("Clean target directory starting with", f.build.targetDir)
			err := os.RemoveAll(f.build.targetDir)
			if err != nil {
				return err
			}
			return os.MkdirAll(f.build.targetDir, os.ModeDir|0700)
		})
	return err
}

func (f *runLocalFlags) makeDistForAdkWebUI() error {
	err := util.LogStartStop("Making dist for Adk Web UI",
		func(p util.Printer) error {
			cmd := exec.Command("ng", "build", "--output-path="+f.build.targetDir)
			cmd.Dir = f.source.webuiDir
			return util.LogCommand(cmd, p)
		})
	return err
}

func (f *runLocalFlags) runLocal() error {
	var err error

	err = f.cleanTemp()
	if err != nil {
		return err
	}
	err = f.makeDistForAdkWebUI()
	if err != nil {
		return err
	}

	return nil
}


Review comment (bug):
nit: you can skip defining the error here and init it in the next line using `:=` operator

Code after:
}

func (f *runLocalFlags) cleanTemp() error {
	return util.LogStartStop("Cleaning target directory",
		func(p util.Printer) error {
			p("Clean target directory starting with", f.build.targetDir)
			err := os.RemoveAll(f.build.targetDir)
			if err != nil {
				return fmt.Errorf("failed to clean target directory %v: %v", f.build.targetDir, err)
			}
			err = os.MkdirAll(f.build.targetDir, os.ModeDir|0700)
			if err != nil {
				return fmt.Errorf("failed to create the target directory %v: %v", f.build.targetDir, err)
			}
			return nil
		})
}

func (f *runLocalFlags) ngBuildADKWebUI() error {
	return util.LogStartStop("Building ADK Web UI",
		func(p util.Printer) error {
			cmd := exec.Command("ng", "build", "--output-path="+f.build.targetDir)
			cmd.Dir = f.source.webuiDir
			return util.LogCommand(cmd, p)
		})
}

func (f *runLocalFlags) buildWebui() error {
	err := f.cleanTemp()
	if err != nil {
		return err
	}
	return f.ngBuildADKWebUI()
}
Repository: google/adk-go
File: cmd/adkgo/root/build/webui/webui.go (Go)

Code before:
			err := os.RemoveAll(f.build.targetDir)
			if err != nil {
				return err
			}
			return os.MkdirAll(f.build.targetDir, os.ModeDir|0700)
		})
	return err
}

func (f *runLocalFlags) makeDistForAdkWebUI() error {
	err := util.LogStartStop("Making dist for Adk Web UI",
		func(p util.Printer) error {
			cmd := exec.Command("ng", "build", "--output-path="+f.build.targetDir)
			cmd.Dir = f.source.webuiDir
			return util.LogCommand(cmd, p)
		})
	return err
}

func (f *runLocalFlags) runLocal() error {
	var err error

	err = f.cleanTemp()
	if err != nil {
		return err
	}
	err = f.makeDistForAdkWebUI()
	if err != nil {
		return err
	}

	return nil
}


Review comment (bug):
Could you add error context here and everywhere? Useful when a function has multiple error exits, with error context it's immediately easy to see the whole error trace.

e.g. `fmt.Errorf("failed to clean up temp: %w", err)` here and everywhere.

Code after:
			p("Clean target directory starting with", f.build.targetDir)
			err := os.RemoveAll(f.build.targetDir)
			if err != nil {
				return fmt.Errorf("failed to clean target directory %v: %v", f.build.targetDir, err)
			}
			err = os.MkdirAll(f.build.targetDir, os.ModeDir|0700)
			if err != nil {
				return fmt.Errorf("failed to create the target directory %v: %v", f.build.targetDir, err)
			}
			return nil
		})
}

func (f *runLocalFlags) ngBuildADKWebUI() error {
	return util.LogStartStop("Building ADK Web UI",
		func(p util.Printer) error {
			cmd := exec.Command("ng", "build", "--output-path="+f.build.targetDir)
			cmd.Dir = f.source.webuiDir
			return util.LogCommand(cmd, p)
		})
}

func (f *runLocalFlags) buildWebui() error {
	err := f.cleanTemp()
	if err != nil {
		return err
	}
	return f.ngBuildADKWebUI()
}
Repository: google/adk-go
File: internal/toolinternal/context.go (Go)

Code before:
}

func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions *session.EventActions) tool.Context {
	if functionCallID == "" {
		functionCallID = uuid.NewString()
	}
	cbCtx := contextinternal.NewCallbackContext(ctx)

	return &toolContext{
		CallbackContext:   cbCtx,
		invocationContext: ctx,
		functionCallID:    functionCallID,
		eventActions:      actions,
		artifacts: &internalArtifacts{
			Artifacts:    ctx.Artifacts(),
			eventActions: actions,
		},
	}
}

type toolContext struct {
	agent.CallbackContext
	invocationContext agent.InvocationContext
	functionCallID    string
	eventActions      *session.EventActions
	artifacts         agent.Artifacts
}

func (c *toolContext) Artifacts() agent.Artifacts {
	return c.artifacts
}

func (c *toolContext) FunctionCallID() string {
	return c.functionCallID
}

func (c *toolContext) Actions() *session.EventActions {
	return c.eventActions
}

func (c *toolContext) AgentName() string {
	return c.invocationContext.Agent().Name()
}

func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) {
	return c.invocationContext.Memory().Search(ctx, query)
}


Review comment (suggestion):
```suggestion
	artifacts         internalArtifacts
```

In case we'd need to access `internalArtifacts`'s fields in toolContext method.

Code after:
}

func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions *session.EventActions) tool.Context {
	if functionCallID == "" {
		functionCallID = uuid.NewString()
	}
	cbCtx := contextinternal.NewCallbackContext(ctx)

	return &toolContext{
		CallbackContext:   cbCtx,
		invocationContext: ctx,
		functionCallID:    functionCallID,
		eventActions:      actions,
		artifacts: &internalArtifacts{
			Artifacts:    ctx.Artifacts(),
			eventActions: actions,
		},
	}
}

type toolContext struct {
	agent.CallbackContext
	invocationContext agent.InvocationContext
	functionCallID    string
	eventActions      *session.EventActions
	artifacts         *internalArtifacts
}

func (c *toolContext) Artifacts() agent.Artifacts {
	return c.artifacts
}

func (c *toolContext) FunctionCallID() string {
	return c.functionCallID
}

func (c *toolContext) Actions() *session.EventActions {
	return c.eventActions
}

func (c *toolContext) AgentName() string {
	return c.invocationContext.Agent().Name()
}

func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) {
	return c.invocationContext.Memory().Search(ctx, query)
}