adding agentCore memory integration in agent execution #4149
Conversation
dhrubo-os
requested review from
HenryL27,
Zhangxunmt,
austintlee,
b4sjoo,
jngz-es,
mingshl,
model-collapse,
pyek-bot,
rbhavna,
sam-herman,
xinyual,
ylwu-amzn and
zane-neo
as code owners
|
Can we remove those emojis in response messages? @dhrubo-os |
Signed-off-by: Dhrubo Saha <[email protected]>
dhrubo-os
force-pushed
the
second_phase
branch
from
c8dee02 to
5fe48b2
Compare
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
dhrubo-os
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #4149 +/- ##
============================================
- Coverage 81.81% 81.19% -0.62%
- Complexity 8847 8970 +123
============================================
Files 761 766 +5
Lines 38099 39021 +922
Branches 4250 4378 +128
============================================
+ Hits 31170 31683 +513
- Misses 5109 5473 +364
- Partials 1820 1865 +45
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
Signed-off-by: Dhrubo Saha <[email protected]>
dhrubo-os
force-pushed
the
second_phase
branch
from
5e42789 to
a7ac7bf
Compare
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
There was a problem hiding this comment.
Thanks for the changes @dhrubo-os ! This is a considerable PR.
A couple things would help to make the review easier:
- A class diagram of the Adapter pattern for the BedrockAgentCoreMemory
- Parallels between existing memory components like memory, message, interaction with AgentCore.
- Memory management between the planner/executor. Currently, they are separate memories. Wondering if we are merging them with agentCore.
I've reviewed most of the PR and left quite a few refactor/nitpick type of comments. Feel free to address them later. However, a few comments are about the functionality and the expected result. Please do take a look at them!
| Map<String, String> parameters = StringUtils.getParameterMap(parser.map()); | ||
| inputDataset = new RemoteInferenceInputDataSet(parameters); | ||
| break; | ||
| case "memory": |
There was a problem hiding this comment.
Use constants (leaving nitpicks so we can track them, they can be addressed later)
| Map<String, String> inputParameters = gson | ||
| .fromJson(input, TypeToken.getParameterized(Map.class, String.class, String.class).getType()); | ||
| extractedParameters.putAll(inputParameters); |
There was a problem hiding this comment.
can we use this please:
Map<String, Object> parsedInputParameters = gson.fromJson(input, TypeToken.getParameterized(Map.class, String.class, Object.class).getType());
extractedParameters.putAll(StringUtils.getParameterMap(parsedInputParameters));
There was a problem hiding this comment.
let's keep this method as in the above PR if possible, thanks
| extractedParameters.putAll(inputParameters); | ||
|
|
||
| // Check if input is a JSON object or array | ||
| String trimmedInput = input.trim(); |
There was a problem hiding this comment.
is there a reason for these changes? i would avoid parameter parsing changes unless 100% required for memory related changes
| input.setMemory(memoryMap); | ||
|
|
||
| assertNotNull("Memory should not be null after setting", input.getMemory()); | ||
| assertEquals("bedrock_agentcore_memory", input.getMemory().get("type")); | ||
| assertEquals("test-arn", input.getMemory().get("memory_arn")); | ||
| assertEquals("us-east-1", input.getMemory().get("region")); |
There was a problem hiding this comment.
this test is only testing the setter, not sure how useful that is
| AgentMLInput agentMLInput = (AgentMLInput) input; | ||
|
|
||
| // DEBUG: Log agentMLInput memory field immediately after cast | ||
| log.info("DEBUG: AgentMLInput memory field immediately after cast: {}", agentMLInput.getMemory()); |
There was a problem hiding this comment.
can this be removed? info log
| if (executorMemoryId != null) { | ||
| reactParams.put("executor_memory_id", executorMemoryId); | ||
| allParams.put("executor_memory_id", executorMemoryId); // Track executor memory ID | ||
| } |
There was a problem hiding this comment.
for reactParams we should use key memory_id so that we don't need custom logic in chat agent for when it is being used as executor, chat agent should not know this information as it is an individual agent, just the caller is not human
| reactParams.put("memory_session_token", allParams.get("memory_session_token")); | ||
|
|
||
| // CRITICAL FIX: Pass executor memory ID to executor agents | ||
| if (executorMemoryId != null) { |
There was a problem hiding this comment.
also this param seems specific to agentcore, maybe consider rename
| } | ||
|
|
||
| // CRITICAL FIX: Use executor memory ID for executor agent calls to maintain shared conversation context | ||
| if (executorMemoryId != null) { |
There was a problem hiding this comment.
same check as above, is this intentional?
|
|
||
| // CRITICAL FIX: Use executor memory ID for executor agent calls to maintain shared conversation context | ||
| if (executorMemoryId != null) { | ||
| reactParams.put("session_id", executorMemoryId); |
There was a problem hiding this comment.
both keys: session_id and executor_memory_id has the same value
| // Last resort: generate new session_id only if none exists | ||
| String sessionId = "bedrock-session-" + System.currentTimeMillis(); | ||
| reactParams.put("session_id", sessionId); | ||
| allParams.put("session_id", sessionId); |
There was a problem hiding this comment.
any reason for not adding session_id to allParams in the above if conditions
| if (memory instanceof ConversationIndexMemory) { | ||
| saveTraceData( | ||
| (ConversationIndexMemory) memory, | ||
| memory.getType(), | ||
| stepToExecute, | ||
| results.get(STEP_RESULT_FIELD), | ||
| conversationId, | ||
| false, | ||
| parentInteractionId, | ||
| traceNumber, | ||
| "PlanExecuteReflect Agent" | ||
| ); | ||
| } |
There was a problem hiding this comment.
are we not storing traceData for agentCore?
| log.info("Skipping interaction save for memory type: {}", memory.getClass().getSimpleName()); | ||
| List<ModelTensors> finalModelTensors = createModelTensors( | ||
| memory.getConversationId(), | ||
| reactAgentMemoryId, |
There was a problem hiding this comment.
reactAgentMemoryId is passed twice, this should be the memory_id of the conversation itself, unsure how agentcore is handling planner/executor memories, is this intentional?
| if (finalResult != null && !finalResult.isEmpty()) { | ||
| // Only add response tensor for non-ConversationIndexMemory cases | ||
| // ConversationIndexMemory handles this in the calling method | ||
| } |
There was a problem hiding this comment.
this condition does nothing
| } else { | ||
| log | ||
| .warn( | ||
| "⚠️ NO RECORDS FOUND: sessionId={}, memoryId={} - This may indicate session mismatch or no events saved yet", |
There was a problem hiding this comment.
Could you please remove those emojis in logs to make it look more product ready?
| public void getMessages(ActionListener<List<Interaction>> listener) { | ||
| log | ||
| .info( | ||
| "🔍 RETRIEVE FOR COMPATIBILITY START: memorySessionId={}, agentId={}, memoryId={}", |
| bedrockClient.listMemoryRecords(getMemoryId(), this.sessionId, this.agentId, ActionListener.wrap(records -> { | ||
| log | ||
| .info( | ||
| "📥 COMPATIBILITY RESPONSE: Found {} records for sessionId={}, memoryId={}", |
| for (BedrockAgentCoreMemoryRecord record : records) { | ||
| log | ||
| .info( | ||
| "📋 CONVERTING RECORD: type={}, content='{}', sessionId={}", |
|
|
||
| @Override | ||
| public void getMessages(String id, ActionListener<BedrockAgentCoreMemoryRecord> listener) { | ||
| log.info("RETRIEVE REQUEST START: requestedSessionId={}, memorySessionId={}, memoryId={}", id, this.sessionId, getMemoryId()); |
There was a problem hiding this comment.
Seems log too detailed, considering downgrade it to DEBUG
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
dhrubo-os
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
3 participants
Description
[adding agentCore memory integration in agent execution]
Related Issues
Resolves #[Issue number to be closed when this PR is merged]
Check List
--signoff.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.