feat(memory): implement Vertex AI Memory Bank service with tests#291
feat(memory): implement Vertex AI Memory Bank service with tests#291AmaadMartin wants to merge 19 commits intogoogle:mainfrom
Conversation
kalenkevich
left a comment
kalenkevich
left a comment
There was a problem hiding this comment.
this is partial review
I will do the rest during the day
Thanks!
| import {Client} from '@google-cloud/vertexai/build/src/genai/client.js'; | ||
| import {Memories} from '@google-cloud/vertexai/build/src/genai/memories.js'; | ||
| import { | ||
| AgentEngineMemoryConfig, | ||
| CreateAgentEngineMemoryRequestParameters, | ||
| GenerateAgentEngineMemoriesConfig, | ||
| GenerateAgentEngineMemoriesRequestParameters, | ||
| GenerateMemoriesRequestDirectContentsSourceEvent, | ||
| MemoryMetadataValue, | ||
| RetrieveAgentEngineMemoriesRequestParameters, | ||
| } from '@google-cloud/vertexai/build/src/genai/types.js'; |
There was a problem hiding this comment.
can we import everything just from @google-cloud/vertexai?
| const retrievedMemoriesResponse = | ||
| await this.memories.retrieveInternal(params); | ||
|
|
||
| logger.info('Search memory response received.'); |
There was a problem hiding this comment.
No need this to be info level, debug is enough.
| await this._addMemoriesViaGenerateDirectMemoriesSource(request); | ||
| return; |
There was a problem hiding this comment.
nit: could be just
return this._addMemoriesViaGenerateDirectMemoriesSource(request);
| const part: Part = {text: retrievedMemory.memory.fact}; | ||
| const content: Content = { | ||
| parts: [part], | ||
| role: 'user', | ||
| }; |
There was a problem hiding this comment.
you can use const content = createUserContent(retrievedMemory.memory.fact) util function from @google/genai
| private _shouldFilterOutEvent(content?: Content): boolean { | ||
| if (!content || !content.parts) { | ||
| return true; | ||
| } | ||
| for (const part of content.parts) { | ||
| const explicitPart: Part = part; | ||
| if ( | ||
| explicitPart.text || | ||
| explicitPart.inlineData || | ||
| explicitPart.fileData | ||
| ) { | ||
| return false; | ||
| } | ||
| } | ||
| return true; | ||
| } |
There was a problem hiding this comment.
Here and below:
this could be just a util function no need to be part of the class.
Check where this. is not called in the methods and make them as util functions. It can happen that after doing that you will see that other methods now also could be util functions as they used this just to call a method which became an util function.
There was a problem hiding this comment.
| private _shouldFilterOutEvent(content?: Content): boolean { | |
| if (!content || !content.parts) { | |
| return true; | |
| } | |
| for (const part of content.parts) { | |
| const explicitPart: Part = part; | |
| if ( | |
| explicitPart.text || | |
| explicitPart.inlineData || | |
| explicitPart.fileData | |
| ) { | |
| return false; | |
| } | |
| } | |
| return true; | |
| } | |
| function shouldFilterOutEvent(content?: Content): boolean { | |
| return !(content?.parts || []).some(p => p.text || p.inlineData || p.fileData); | |
| } |
| if (this._shouldFilterOutEvent(event.content)) { | ||
| continue; | ||
| } | ||
| if (event.content) { |
There was a problem hiding this comment.
this check does not make sence as you just checked that in _shouldFilterOutEvent function
| config: config, | ||
| }; | ||
| const operation = await this.memories.generateInternal(params); | ||
| logger.info('Generate memory response received.'); |
There was a problem hiding this comment.
make it debug as well
| return vertexMetadata; | ||
| } | ||
|
|
||
| private _toVertexMetadataValue( |
There was a problem hiding this comment.
can be just an util function
| const knownFields = [ | ||
| 'disableConsolidation', | ||
| 'waitForCompletion', | ||
| 'revisionLabels', | ||
| 'revisionExpireTime', | ||
| 'revisionTtl', | ||
| 'disableMemoryRevisions', | ||
| 'metadataMergeStrategy', | ||
| 'allowedTopics', | ||
| ]; |
There was a problem hiding this comment.
Here and in all other similar places
move this on the file level as you don't need to create this string list every time with exact same values, you can do it only once.
| interface MemoryEntryWithMetadata extends MemoryEntry { | ||
| customMetadata?: Record<string, unknown>; | ||
| } |
There was a problem hiding this comment.
move this on file level
2 participants
Link to Issue or Description of Change
Or, if no issue exists, describe the change:
Problem:
The
adk-jsrepository lacked an implementation forVertexAiMemoryBankServiceto interact with Vertex AI Memory Bank for storage and retrieval of semantic memories, which was already available in the Python implementation.Solution:
VertexAiMemoryBankService: Createdcore/src/memory/vertex_ai_memory_bank_service.tsleveraging theMemoriesclient from@google-cloud/vertexaito support adding sessions, adding explicit memories, and searching memories.core/test/memory/vertex_ai_memory_bank_service_test.tsachieving 94.78% line coverage.tests/integration/memory/vertex_ai_memory_bank_service_test.tsto verify interaction with theRunnerandLOAD_MEMORYtool.Testing Plan
Unit Tests:
Summary of passed results:
VertexAiMemoryBankServicepass with high coverage (94.78%), covering initialization, adding sessions, adding memories, searching, and metadata conversion.Manual End-to-End (E2E) Tests:
tests/integration/memory/vertex_ai_memory_bank_service_test.tssimulating the full flow with a mock client,Runner, and theLOAD_MEMORYtool.Checklist