fix: update built-in load_memory tool for compatibility with VertexAiSessionService

[nicholasbreckwoldt]

• Closes: Built-in load-memory tool does not work in combination with VertexAiSessionService #792

Problem:
Built-in ADK load_memory tool throws an error when using VertexAiSessionService to manage session event persistence. Downstream, the function call response is passed to structpb.NewStruct() when used with this custom session service implementation. structpb is strictly incompatible with native Go structs. Passing a struct directly causes a runtime error / failure.

Solution:
This PR updates the built-in load_memory tool logic to convert the [memory.searchResponse] struct into a generic map[string]any before returning. This is done using a JSON round-trip (marshall / unmarshal) to ensure more general compatability, including with VertexAiSessionService.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

The following local tests were run and passed successfully. Note minor adjustment to the TestLoadMemoryTool_Run test, i.e. marshalling tool responses into the expected [memory.SearchResponse] type to make handling more robust and type safe.

go test ./tool/loadmemorytool

Test results below.

Manual End-to-End (E2E) Tests:

Manually tested with ADK agent configured with VertexAISessionService and PR-proposed changes. Code snippet and screenshot showing E2E testing below:

...
  // Initialise Vertex AI Engine Session Service
  sessionService, err := vertexai.NewSessionService(ctx, vertexai.VertexAIServiceConfig{
  	ProjectID: "PROJECT",
  	Location:  "LOCATION",
      ReasoningEngine: "REASONING_ENGINE",
  })
  if err != nil {
      log.Fatalf("failed to create session service")
  }

// Create new agent
  adkAgent, err := llmagent.New(llmagent.Config{
      Name:        "agent",
      Model:       model,
      Description: "Description",
      Instruction: "You are a helpful assistant..",
      Tools: []tool.Tool{
          loadmemorytool.New(),
      },
  })
  if err != nil {
      log.Fatalf("failed to create agent: %v", err)
  }  

  // Configure launcher
  config := &launcher.Config{
      AgentLoader: agent.NewSingleLoader(adkAgent),
      SessionService: sessionService,
      MemoryService: memory.InMemoryService(),
  }

...

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

[kdroste-google]

Hi, @nicholasbreckwoldt , thank you for your contribution!.

The underlying problem in that case seems to be in session/vertexai/vertexai_client.go createAiplatformpbContent, which calls
response, err := structpb.NewStruct(part.FunctionResponse.Response)
when part.FunctionResponse.Response is map[string]interface {} ["memories": []google.golang.org/adk/memory.Entry len: 0, cap: 0, [], ]

I would say that the fix should go here.

[nicholasbreckwoldt]

@kdroste-google Thank you for the feedback. Agree that source of issue does indeed originate in the session/vertexai/vertexai_client.go.

My original thinking was to instead incorporate the fix at the load_memory tool itself to avoid the JSON round-trip overhead for each function call response. However, fixing at source as you point out (i.e. vertexai_client) would likely be more robust in preventing the same error being created across other toolsets.

I have migrated the fix across to there and re-run Manual End-to-End (E2E) testing:

Look forward to your review.

[nicholasbreckwoldt]

@kdroste-google Thanks for pointing this out and agreed.

Based on what I identified, these are the potential error prone conversion of Go types to struct protobuf type instances (all in the adk-go/session/vertexai/vertexai_client.go file):

• func createSession: creating new session with initial non-empty state
• func appendEvent: applying event.Actions.StateDelta
• func createAiplatformpbContent: parts.FunctionalCall.Args, parts.FunctionResponse.Response conversion
• func createAiplatformpbMetadata: event.CustomMetadata conversion

Given that these all occur within the adk-go/session/vertexai package, my initial thinking would be to include the above suggested util function within this same package alongside other existing utils already there.

Do you foresee any reason to instead place this util function elsewhere (for example the adk-go/internal/utils package) for broader use down the line?

[kdroste-google]

At this moment it's not needed elsewhere, so let's keep it local and unexported. We can always move it to a shared package later on.

[kdroste-google]

Could you please add some tests?

[nicholasbreckwoldt]

@kdroste-google Have simplified and updated function naming as suggested, and also extended test coverage. Test coverage extension focused on validating that new toStructPB util behaves appropriately and caters for richer session state deltas (Go structs, etc). Also covered createAiplatformpbMetadata and createAiplatformpbContent cases where conversion to structpb is applicable.

Look forward to your feedback.

[kdroste-google]

Looks good, thank you!

[nicholasbreckwoldt]

@kdroste-google I don't appear to have the ability to merge after your approval. Would you kindly merge on my behalf if no further steps required?

[kdroste-google]

No problem! :)