Support messages array in all memory types + chat history in AGUI agent#4645
Conversation
jiapingzeng
requested review from
HenryL27,
Zhangxunmt,
austintlee,
b4sjoo,
dhrubo-os,
jngz-es,
mingshl,
model-collapse,
pyek-bot,
rbhavna,
sam-herman,
xinyual,
ylwu-amzn and
zane-neo
as code owners
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR introduces structured message parsing across multiple LLM providers (Bedrock, OpenAI, Gemini), extends the Memory interface with structured message handling, refactors agent execution to use input messages directly instead of storing in memory, and updates memory implementations and agent runners to support unified message interfaces. Changes
Sequence DiagramsequenceDiagram
participant Agent as Agent Executor
participant Provider as Model Provider
participant LLM as LLM Service
participant Memory as Memory System
participant Runner as Agent Runner
Agent->>Agent: Extract input messages from input
Agent->>Runner: setInputMessages(inputMessages)
Runner->>Memory: getStructuredMessages()
Memory-->>Runner: historical messages + context
Runner->>LLM: Send prompt with history
LLM-->>Provider: Return JSON response
Provider->>Provider: parseResponseMessage(json)
Provider->>Provider: Extract text/tools from response
Provider-->>Runner: Message with ContentBlocks + ToolCalls
Runner->>Runner: Process tool interactions
Runner->>Memory: saveStructuredMessages(newMessages)
Memory-->>Runner: Acknowledgment
Runner-->>Agent: Final answer
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
jiapingzeng
temporarily deployed
to
ml-commons-cicd-env-require-approval
GitHub Actions
Inactive
jiapingzeng
temporarily deployed
to
ml-commons-cicd-env-require-approval
GitHub Actions
Inactive
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java`:
- Around line 1319-1417: The method extractMessagePairs currently assumes
messages is non-null and will NPE; add an early guard at the start of
extractMessagePairs to return Collections.emptyList() (or new ArrayList<>())
when messages is null or messages.isEmpty(), ensuring all callers receive an
empty list instead of throwing; update the beginning of the
extractMessagePairs(List<Message> messages, String sessionId, String appType)
method accordingly.
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java`:
- Around line 271-303: The unified-interface branch (usesUnifiedInterface in
MLChatAgentRunner) assumes inputMessages is non-null and that
ModelProviderFactory.getProvider(mlAgent.getModel().getModelProvider()) returns
a valid provider; add quick validation before calling
memory.getStructuredMessages: check inputMessages != null and not empty, verify
mlAgent.getModel() and mlAgent.getModel().getModelProvider() are present, call
ModelProviderFactory.getProvider(...) and fail-fast with listener.onFailure(...)
if it returns null or throws; only then proceed to
memory.getStructuredMessages(...) / memory.saveStructuredMessages(...) and
runAgent(...).
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/AgenticConversationMemory.java`:
- Around line 429-502: saveStructuredMessages issues concurrent writes which can
arrive out of order despite using messageId=i; to fix, make saves sequential or
assign a strictly monotonic sequence number instead of relying on loop index.
Modify saveStructuredMessages to either: 1) replace the parallel client.execute
loop with a recursive/iterative sequential publisher (e.g., a private helper
saveNext(int index) that builds MLAddMemoriesInput/MLAddMemoriesRequest and
calls client.execute and on success calls saveNext(index+1), propagating
failures to listener), or 2) keep parallel writes but generate and persist a
monotonic sequence number for MLAddMemoriesInput.messageId (e.g., from an
AtomicLong sequence obtained from the memory container) so ordering is
guaranteed; adjust handling around MLAddMemoriesInput, MLAddMemoriesRequest,
client.execute, remaining/hasError, and messageId accordingly.
🧹 Nitpick comments (2)
common/src/main/java/org/opensearch/ml/common/memory/Memory.java (1)
21-24: Consider makingMAX_MESSAGES_TO_RETRIEVEconfigurable.The hardcoded value of 10,000 messages could lead to memory pressure or performance issues when retrieving large conversation histories. Consider making this configurable via settings, or at minimum, documenting the rationale for this specific value.
ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemory.java (1)
515-637: Preserve the first failure when batching structured-message saves.If an early save fails but the last succeeds, callers only see a generic error. Capturing the first exception keeps the root cause.
♻️ Suggested adjustment
AtomicInteger remaining = new AtomicInteger(messages.size()); AtomicBoolean hasError = new AtomicBoolean(false); + java.util.concurrent.atomic.AtomicReference<Exception> firstError = new java.util.concurrent.atomic.AtomicReference<>(); @@ if (remaining.decrementAndGet() == 0) { if (hasError.get()) { - listener.onFailure(new RuntimeException("One or more structured messages failed to save")); + Exception err = firstError.get(); + listener.onFailure(err != null ? err : new RuntimeException("One or more structured messages failed to save")); } else { listener.onResponse(null); } } }, e -> { log.error("Failed to save structured message {} of {} to remote session {}", index + 1, messages.size(), conversationId, e); hasError.set(true); + firstError.compareAndSet(null, e); if (remaining.decrementAndGet() == 0) { - listener.onFailure(e); + Exception err = firstError.get(); + listener.onFailure(err != null ? err : e); } }));
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (26)
common/src/main/java/org/opensearch/ml/common/agent/BedrockConverseModelProvider.javacommon/src/main/java/org/opensearch/ml/common/agent/GeminiV1BetaGenerateContentModelProvider.javacommon/src/main/java/org/opensearch/ml/common/agent/OpenaiV1ChatCompletionsModelProvider.javacommon/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.javacommon/src/main/java/org/opensearch/ml/common/input/execute/agent/ContentBlock.javacommon/src/main/java/org/opensearch/ml/common/input/execute/agent/Message.javacommon/src/main/java/org/opensearch/ml/common/memory/Memory.javacommon/src/main/java/org/opensearch/ml/common/model/ModelProvider.javacommon/src/test/java/org/opensearch/ml/common/agent/BedrockConverseModelProviderTest.javacommon/src/test/java/org/opensearch/ml/common/agent/OpenaiV1ChatCompletionsModelProviderTest.javacommon/src/test/java/org/opensearch/ml/common/agui/AGUIInputConverterTest.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentRunner.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.javaml-algorithms/src/main/java/org/opensearch/ml/engine/memory/AgenticConversationMemory.javaml-algorithms/src/main/java/org/opensearch/ml/engine/memory/ConversationIndexMemory.javaml-algorithms/src/main/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemory.javaml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/AgentUtilsTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandlerTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/memory/AgenticConversationMemoryTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/memory/ConversationIndexMemoryTest.javaml-algorithms/src/test/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemoryTest.java
💤 Files with no reviewable changes (1)
- ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2026-01-13T04:54:54.656Z
Learnt from: akolarkunnu
Repo: opensearch-project/ml-commons PR: 3919
File: ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java:394-409
Timestamp: 2026-01-13T04:54:54.656Z
Learning: In ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java, the test encrypt_ThrowExceptionWhenInitMLConfigIndex uses exceptionRule (not listener-based error assertions) because it mocks mlIndicesHandler.initMLConfigIndex() with doThrow() to throw synchronously before any listener callbacks occur, validating sync exception propagation rather than async listener error delivery.
Applied to files:
ml-algorithms/src/test/java/org/opensearch/ml/engine/memory/ConversationIndexMemoryTest.javaml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java
🧬 Code graph analysis (6)
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentRunner.java (1)
search-processors/src/main/java/org/opensearch/searchpipelines/questionanswering/generative/prompt/PromptUtil.java (1)
Message(540-593)
common/src/main/java/org/opensearch/ml/common/agent/OpenaiV1ChatCompletionsModelProvider.java (1)
common/src/main/java/org/opensearch/ml/common/model/ModelProvider.java (1)
ModelProvider(22-114)
ml-algorithms/src/test/java/org/opensearch/ml/engine/memory/AgenticConversationMemoryTest.java (2)
common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/memory/MLAddMemoriesAction.java (1)
MLAddMemoriesAction(10-17)common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/memory/MLSearchMemoriesAction.java (1)
MLSearchMemoriesAction(11-18)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java (2)
common/src/main/java/org/opensearch/ml/common/input/execute/agent/AgentMLInput.java (1)
org(32-178)search-processors/src/main/java/org/opensearch/searchpipelines/questionanswering/generative/prompt/PromptUtil.java (1)
Message(540-593)
common/src/main/java/org/opensearch/ml/common/agent/BedrockConverseModelProvider.java (1)
common/src/main/java/org/opensearch/ml/common/input/execute/agent/AgentMLInput.java (1)
org(32-178)
ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/AgenticConversationMemory.java (2)
common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/memory/MLSearchMemoriesAction.java (1)
MLSearchMemoriesAction(11-18)common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/memory/MLAddMemoriesAction.java (1)
MLAddMemoriesAction(10-17)
🔇 Additional comments (52)
common/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.java (1)
112-113: LGTM!The propagation of
threadIdasmemory_idis appropriate for enabling memory-scoped processing in AGUI flows. The inline comment clearly documents the purpose.common/src/main/java/org/opensearch/ml/common/input/execute/agent/ContentBlock.java (1)
8-9: LGTM!Adding
@JsonInclude(JsonInclude.Include.NON_NULL)is appropriate for producing cleaner JSON output by excluding null fields. This aligns with the same annotation applied toMessage.java.Also applies to: 18-18
common/src/main/java/org/opensearch/ml/common/input/execute/agent/Message.java (1)
10-11: LGTM!The Jackson annotation for excluding null fields is consistent with
ContentBlock.java, ensuring uniform serialization behavior across the message-related DTOs.Also applies to: 20-20
common/src/test/java/org/opensearch/ml/common/agui/AGUIInputConverterTest.java (2)
687-698: LGTM!The test properly validates the new
memory_idparameter propagation fromthreadId. It correctly asserts both thememory_idandAGUI_PARAM_THREAD_IDvalues are present and equal.
700-717: LGTM!The helper method
buildMinimalAGUIInputis well-structured, creating a valid AGUI input with all required fields. It can be reused for additional test cases if needed.common/src/main/java/org/opensearch/ml/common/model/ModelProvider.java (1)
77-83: LGTM!The new abstract method
parseResponseMessageprovides a clean extension to theModelProvidercontract, enabling unified response parsing across different LLM providers. The Javadoc clearly documents the expected behavior.common/src/main/java/org/opensearch/ml/common/memory/Memory.java (1)
46-66: LGTM!The default implementations using
UnsupportedOperationExceptionprovide clean backward compatibility, allowing existing memory implementations to remain unchanged while new implementations can override these methods. The use of fully qualifiedorg.opensearch.ml.common.input.execute.agent.Messageavoids potential naming conflicts with the generic type parameterT extends Message.ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java (3)
514-519: LGTM!The extraction of
inputMessagesfromAgentMLInputwhen the input type isMESSAGESis correctly implemented. The conditional check ensures thatinputMessagesis only populated when appropriate, and the@SuppressWarnings("unchecked")annotation is justified for the type cast.
434-436: LGTM!Passing
nullforinputMessagesin the existing memory path is correct since these code paths handle scenarios where messages are retrieved from memory rather than passed directly.Also applies to: 458-460
732-732: No action required—null handling is already implemented.The
inputMessagesparameter can be null (when input type is notMESSAGES), and the code safely handles this case. The interface documentation explicitly states the parameter "may be null." All downstream consumers, includingConversationIndexMemory.saveStructuredMessages(),RemoteAgenticConversationMemory.saveStructuredMessages(), andAgenticConversationMemory.saveStructuredMessages(), have explicit null checks that return successfully when messages are null or empty. This behavior is also covered by existing unit tests.common/src/main/java/org/opensearch/ml/common/agent/OpenaiV1ChatCompletionsModelProvider.java (2)
8-8: LGTM!The new imports are correctly added to support the
parseResponseMessageimplementation.Also applies to: 25-25, 32-32
356-409: Implementation looks correct with a minor null-handling note.The
parseResponseMessagemethod correctly handles OpenAI's response format including text content, tool calls, and tool result messages. The default role of "assistant" is appropriate for response parsing.One minor consideration: on line 381,
String.valueOf(tc.get("id"))will return the literal string"null"ifidis null, rather than an empty string. While the checktc.get("id") != nullis present, consider using the same pattern as line 386 withgetOrDefaultfor consistency:💡 Optional consistency improvement
- String id = tc.get("id") != null ? String.valueOf(tc.get("id")) : ""; + String id = tc.get("id") != null ? String.valueOf(tc.get("id")) : "";The current implementation is already correct - the ternary handles the null case. This is just a note for awareness.
common/src/test/java/org/opensearch/ml/common/agent/OpenaiV1ChatCompletionsModelProviderTest.java (1)
1104-1184: Good coverage for parseResponseMessage edge cases.These tests exercise text-only, tool calls, tool results, null/empty content, and default role paths, which should catch most regressions in the parser.
ml-algorithms/src/test/java/org/opensearch/ml/engine/memory/AgenticConversationMemoryTest.java (1)
777-1041: Solid structured-message coverage in memory tests.The new cases cover null/empty input handling, success/failure paths, and container-ID validation, which should harden the new APIs.
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentRunner.java (2)
8-13: Import-only change.
40-44: Default setter keeps backward compatibility.A no-op default preserves existing implementations while allowing runners to opt into structured inputs.
common/src/main/java/org/opensearch/ml/common/agent/BedrockConverseModelProvider.java (2)
8-35: Import-only change.
412-485: Bedrock response parsing looks solid.The method cleanly handles text blocks, tool calls, tool results, and role defaults with appropriate null checks.
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/AgentUtilsTest.java (1)
2061-2257: Great coverage for structured message extraction utilities.These cases cover null/empty content, mixed block types, whitespace handling, and message pairing behavior.
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java (1)
1287-1307: Text-block extraction looks solid.Concise helper and trimming behavior keep downstream pairing predictable.
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java (1)
1596-1651: Good coverage forsetInputMessagesbehavior.Covers stored, null, and default no-op paths without over-mocking.
ml-algorithms/src/test/java/org/opensearch/ml/engine/memory/ConversationIndexMemoryTest.java (1)
162-304: Structured-message tests hit the important cases.Null/empty inputs, partial failures, and conversion paths are exercised.
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.java (2)
130-325: AG‑UI streaming aggregation fits the final-answer flow.Accumulating content and emitting run-finished plus final answer keeps the AG‑UI stream consistent.
436-455: Final-answer wrapper matches Bedrock response shape.Using
stopReason = end_turnand a content array aligns with downstream parsing expectations.ml-algorithms/src/test/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemoryTest.java (1)
334-395: Structured-message edge cases are well covered.Missing container-id, null, and empty inputs are handled.
ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemory.java (1)
74-78: Deterministic ordering key is a good addition.Sorting with
message_idalongside created_time should stabilize retrieval.common/src/test/java/org/opensearch/ml/common/agent/BedrockConverseModelProviderTest.java (1)
1006-1088: Solid coverage for parseResponseMessage scenarios.
These cases should keep text/tool parsing behavior stable across edge inputs.ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandlerTest.java (2)
9-18: Import additions align with new test usage.
242-290: Nice coverage for final-answer response creation.
Covers non-null, null, and empty text paths.ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java (11)
73-90: Import additions are consistent with the structured-message flow.
167-207: Input message injection looks straightforward.
230-232: Early exit path is unchanged in behavior.
412-425: AG-UI tool-result routing looks consistent with the unified-tools flow.
429-443: Unified-interface context wiring in runReAct looks consistent.
519-523: Call-site propagation of unified-interface context looks correct.Also applies to: 572-575, 717-720
955-1027: Final-answer persistence now cleanly branches for unified vs legacy.
1034-1084: Structured assistant message save helper is clear and self-contained.
1232-1236: Max-iteration paths now consistently pass unified-interface context.Also applies to: 1293-1297
1482-1491: Unified-tools wiring across backend/frontend is consistent.Also applies to: 1540-1541
1594-1598: AG-UI tool-result reinjection reuses the unified-tools path correctly.ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/ConversationIndexMemory.java (5)
11-26: Import additions are aligned with structured-message support.
36-36: Memory type parameterization update looks consistent.
80-104: save(...) signature updates are consistent with the new memory Message type.
121-164: Structured message reconstruction from interactions looks solid.
166-233: Sequential save path with aggregated error handling is robust.ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/AgenticConversationMemory.java (7)
8-55: Import updates align with new structured-message functionality.
65-73: Class signature and metadata field additions look consistent.
95-173: save(...) updates are consistent with memory-container storage format.
227-227: Updated timestamp merge is correct.
267-375: Interaction parsing updates look consistent with the new schema.
382-426: Structured-message retrieval path looks clean.
590-592: Trace timestamp handling looks consistent.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| @Override | ||
| public Message parseResponseMessage(String json) { | ||
| throw new UnsupportedOperationException("parseResponseMessage is not yet supported for Gemini model provider"); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate all parseResponseMessage call sites and Gemini provider usage
rg -n "parseResponseMessage\(" --glob "*.java" -A 2 -B 2Repository: opensearch-project/ml-commons
Length of output: 13816
🏁 Script executed:
#!/bin/bash
# Find all usages of GeminiV1BetaGenerateContentModelProvider
rg -n "GeminiV1BetaGenerateContentModelProvider" --glob "*.java"Repository: opensearch-project/ml-commons
Length of output: 1096
🏁 Script executed:
#!/bin/bash
# Check the parent class/interface to understand the contract
rg -n "class GeminiV1BetaGenerateContentModelProvider" --glob "*.java" -A 5Repository: opensearch-project/ml-commons
Length of output: 1928
🏁 Script executed:
#!/bin/bash
# Find other model providers to see how they implement parseResponseMessage
rg -n "public Message parseResponseMessage" --glob "*.java" -B 5 -A 3Repository: opensearch-project/ml-commons
Length of output: 3424
🏁 Script executed:
#!/bin/bash
# Examine the context around parseResponseMessage call in MLChatAgentRunner
cat -n ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java | sed -n '1050,1075p'Repository: opensearch-project/ml-commons
Length of output: 1395
🏁 Script executed:
#!/bin/bash
# Look for toolInteractions generation and any Gemini guards
rg -n "toolInteractions" --glob "*.java" -B 5 -A 5Repository: opensearch-project/ml-commons
Length of output: 27193
🏁 Script executed:
#!/bin/bash
# Search for any Gemini-specific guards in MLChatAgentRunner
rg -n "GEMINI|Gemini" ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java -B 3 -A 3Repository: opensearch-project/ml-commons
Length of output: 55
🏁 Script executed:
#!/bin/bash
# Check Gemini implementation for tool support and unified interface support
cat -n common/src/main/java/org/opensearch/ml/common/agent/GeminiV1BetaGenerateContentModelProvider.java | head -100Repository: opensearch-project/ml-commons
Length of output: 5598
🏁 Script executed:
#!/bin/bash
# Look for usesUnifiedInterface configuration and Gemini usage
rg -n "usesUnifiedInterface" --glob "*.java" -B 2 -A 2Repository: opensearch-project/ml-commons
Length of output: 8323
🏁 Script executed:
#!/bin/bash
# Check if there are any tests that use Gemini with tool interactions
rg -n "GeminiV1BetaGenerateContentModelProviderTest" --glob "*.java" -A 50 | head -100Repository: opensearch-project/ml-commons
Length of output: 7254
🏁 Script executed:
#!/bin/bash
# Check if Gemini has tool/function calling support
rg -n "tool" common/src/main/java/org/opensearch/ml/common/agent/GeminiV1BetaGenerateContentModelProvider.java -iRepository: opensearch-project/ml-commons
Length of output: 286
🏁 Script executed:
#!/bin/bash
# Check the full Gemini implementation for unified interface support
wc -l common/src/main/java/org/opensearch/ml/common/agent/GeminiV1BetaGenerateContentModelProvider.javaRepository: opensearch-project/ml-commons
Length of output: 174
🏁 Script executed:
#!/bin/bash
# Get full Gemini implementation to understand tool support
sed -n '200,350p' common/src/main/java/org/opensearch/ml/common/agent/GeminiV1BetaGenerateContentModelProvider.javaRepository: opensearch-project/ml-commons
Length of output: 5495
Implement parseResponseMessage for Gemini to avoid losing tool interaction data.
The method throws UnsupportedOperationException, but since the call is wrapped in a try-catch at MLChatAgentRunner:1061, it will not crash the agent. However, when tool interactions exist for Gemini, they will be silently skipped and not stored as structured messages in the unified interface flow. Gemini supports tools (tool configurations are included in its request template), so the parser is needed to properly capture tool call and result messages in conversation history.
| public static List<ConversationIndexMessage> extractMessagePairs(List<Message> messages, String sessionId, String appType) { | ||
| List<ConversationIndexMessage> messagePairs = new ArrayList<>(); | ||
|
|
||
| StringBuilder userTextBuilder = new StringBuilder(); | ||
| StringBuilder assistantTextBuilder = new StringBuilder(); | ||
| boolean skippingTrailingUsers = true; | ||
| String currentRole = null; | ||
|
|
||
| for (int i = messages.size() - 1; i >= 0; i--) { | ||
| Message message = messages.get(i); | ||
|
|
||
| if (message == null || message.getRole() == null) { | ||
| continue; | ||
| } | ||
|
|
||
| String role = message.getRole().toLowerCase(); | ||
|
|
||
| // Skip non-user/assistant roles | ||
| if (!role.equals("user") && !role.equals("assistant")) { | ||
| continue; | ||
| } | ||
|
|
||
| // Skip trailing user messages | ||
| if (skippingTrailingUsers && role.equals("user")) { | ||
| continue; | ||
| } | ||
|
|
||
| if (skippingTrailingUsers && role.equals("assistant")) { | ||
| skippingTrailingUsers = false; | ||
| } | ||
|
|
||
| // Detect role change from user to assistant (going backwards) | ||
| if (currentRole != null && currentRole.equals("user") && role.equals("assistant")) { | ||
| // Save the accumulated pair | ||
| String userText = userTextBuilder.toString().trim(); | ||
| String assistantText = assistantTextBuilder.toString().trim(); | ||
|
|
||
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | ||
| ConversationIndexMessage msg = ConversationIndexMessage | ||
| .conversationIndexMessageBuilder() | ||
| .type(appType) | ||
| .question(userText) | ||
| .response(assistantText) | ||
| .finalAnswer(true) | ||
| .sessionId(sessionId) | ||
| .build(); | ||
|
|
||
| messagePairs.add(msg); | ||
| } | ||
|
|
||
| // Clear buffers for next pair | ||
| userTextBuilder.setLength(0); | ||
| assistantTextBuilder.setLength(0); | ||
| } | ||
|
|
||
| // Extract text | ||
| String text = extractTextFromMessage(message); | ||
|
|
||
| // Accumulate text based on role (prepending since we're going backwards) | ||
| if (role.equals("user")) { | ||
| if (!text.isEmpty()) { | ||
| if (userTextBuilder.length() > 0) { | ||
| userTextBuilder.insert(0, "\n"); | ||
| } | ||
| userTextBuilder.insert(0, text); | ||
| } | ||
| } else if (role.equals("assistant")) { | ||
| if (!text.isEmpty()) { | ||
| if (assistantTextBuilder.length() > 0) { | ||
| assistantTextBuilder.insert(0, "\n"); | ||
| } | ||
| assistantTextBuilder.insert(0, text); | ||
| } | ||
| } | ||
|
|
||
| currentRole = role; | ||
| } | ||
|
|
||
| // Save any remaining pair | ||
| String userText = userTextBuilder.toString().trim(); | ||
| String assistantText = assistantTextBuilder.toString().trim(); | ||
|
|
||
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | ||
| ConversationIndexMessage msg = ConversationIndexMessage | ||
| .conversationIndexMessageBuilder() | ||
| .type(appType) | ||
| .question(userText) | ||
| .response(assistantText) | ||
| .finalAnswer(true) | ||
| .sessionId(sessionId) | ||
| .build(); | ||
|
|
||
| messagePairs.add(msg); | ||
| } | ||
|
|
||
| // Reverse to maintain chronological order | ||
| Collections.reverse(messagePairs); | ||
|
|
||
| return messagePairs; |
There was a problem hiding this comment.
Guard against null/empty message lists.
extractMessagePairs will throw if messages is null; returning an empty list keeps it safe for callers.
🛡️ Suggested guard
public static List<ConversationIndexMessage> extractMessagePairs(List<Message> messages, String sessionId, String appType) {
+ if (messages == null || messages.isEmpty()) {
+ return Collections.emptyList();
+ }
List<ConversationIndexMessage> messagePairs = new ArrayList<>();📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| public static List<ConversationIndexMessage> extractMessagePairs(List<Message> messages, String sessionId, String appType) { | |
| List<ConversationIndexMessage> messagePairs = new ArrayList<>(); | |
| StringBuilder userTextBuilder = new StringBuilder(); | |
| StringBuilder assistantTextBuilder = new StringBuilder(); | |
| boolean skippingTrailingUsers = true; | |
| String currentRole = null; | |
| for (int i = messages.size() - 1; i >= 0; i--) { | |
| Message message = messages.get(i); | |
| if (message == null || message.getRole() == null) { | |
| continue; | |
| } | |
| String role = message.getRole().toLowerCase(); | |
| // Skip non-user/assistant roles | |
| if (!role.equals("user") && !role.equals("assistant")) { | |
| continue; | |
| } | |
| // Skip trailing user messages | |
| if (skippingTrailingUsers && role.equals("user")) { | |
| continue; | |
| } | |
| if (skippingTrailingUsers && role.equals("assistant")) { | |
| skippingTrailingUsers = false; | |
| } | |
| // Detect role change from user to assistant (going backwards) | |
| if (currentRole != null && currentRole.equals("user") && role.equals("assistant")) { | |
| // Save the accumulated pair | |
| String userText = userTextBuilder.toString().trim(); | |
| String assistantText = assistantTextBuilder.toString().trim(); | |
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | |
| ConversationIndexMessage msg = ConversationIndexMessage | |
| .conversationIndexMessageBuilder() | |
| .type(appType) | |
| .question(userText) | |
| .response(assistantText) | |
| .finalAnswer(true) | |
| .sessionId(sessionId) | |
| .build(); | |
| messagePairs.add(msg); | |
| } | |
| // Clear buffers for next pair | |
| userTextBuilder.setLength(0); | |
| assistantTextBuilder.setLength(0); | |
| } | |
| // Extract text | |
| String text = extractTextFromMessage(message); | |
| // Accumulate text based on role (prepending since we're going backwards) | |
| if (role.equals("user")) { | |
| if (!text.isEmpty()) { | |
| if (userTextBuilder.length() > 0) { | |
| userTextBuilder.insert(0, "\n"); | |
| } | |
| userTextBuilder.insert(0, text); | |
| } | |
| } else if (role.equals("assistant")) { | |
| if (!text.isEmpty()) { | |
| if (assistantTextBuilder.length() > 0) { | |
| assistantTextBuilder.insert(0, "\n"); | |
| } | |
| assistantTextBuilder.insert(0, text); | |
| } | |
| } | |
| currentRole = role; | |
| } | |
| // Save any remaining pair | |
| String userText = userTextBuilder.toString().trim(); | |
| String assistantText = assistantTextBuilder.toString().trim(); | |
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | |
| ConversationIndexMessage msg = ConversationIndexMessage | |
| .conversationIndexMessageBuilder() | |
| .type(appType) | |
| .question(userText) | |
| .response(assistantText) | |
| .finalAnswer(true) | |
| .sessionId(sessionId) | |
| .build(); | |
| messagePairs.add(msg); | |
| } | |
| // Reverse to maintain chronological order | |
| Collections.reverse(messagePairs); | |
| return messagePairs; | |
| public static List<ConversationIndexMessage> extractMessagePairs(List<Message> messages, String sessionId, String appType) { | |
| if (messages == null || messages.isEmpty()) { | |
| return Collections.emptyList(); | |
| } | |
| List<ConversationIndexMessage> messagePairs = new ArrayList<>(); | |
| StringBuilder userTextBuilder = new StringBuilder(); | |
| StringBuilder assistantTextBuilder = new StringBuilder(); | |
| boolean skippingTrailingUsers = true; | |
| String currentRole = null; | |
| for (int i = messages.size() - 1; i >= 0; i--) { | |
| Message message = messages.get(i); | |
| if (message == null || message.getRole() == null) { | |
| continue; | |
| } | |
| String role = message.getRole().toLowerCase(); | |
| // Skip non-user/assistant roles | |
| if (!role.equals("user") && !role.equals("assistant")) { | |
| continue; | |
| } | |
| // Skip trailing user messages | |
| if (skippingTrailingUsers && role.equals("user")) { | |
| continue; | |
| } | |
| if (skippingTrailingUsers && role.equals("assistant")) { | |
| skippingTrailingUsers = false; | |
| } | |
| // Detect role change from user to assistant (going backwards) | |
| if (currentRole != null && currentRole.equals("user") && role.equals("assistant")) { | |
| // Save the accumulated pair | |
| String userText = userTextBuilder.toString().trim(); | |
| String assistantText = assistantTextBuilder.toString().trim(); | |
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | |
| ConversationIndexMessage msg = ConversationIndexMessage | |
| .conversationIndexMessageBuilder() | |
| .type(appType) | |
| .question(userText) | |
| .response(assistantText) | |
| .finalAnswer(true) | |
| .sessionId(sessionId) | |
| .build(); | |
| messagePairs.add(msg); | |
| } | |
| // Clear buffers for next pair | |
| userTextBuilder.setLength(0); | |
| assistantTextBuilder.setLength(0); | |
| } | |
| // Extract text | |
| String text = extractTextFromMessage(message); | |
| // Accumulate text based on role (prepending since we're going backwards) | |
| if (role.equals("user")) { | |
| if (!text.isEmpty()) { | |
| if (userTextBuilder.length() > 0) { | |
| userTextBuilder.insert(0, "\n"); | |
| } | |
| userTextBuilder.insert(0, text); | |
| } | |
| } else if (role.equals("assistant")) { | |
| if (!text.isEmpty()) { | |
| if (assistantTextBuilder.length() > 0) { | |
| assistantTextBuilder.insert(0, "\n"); | |
| } | |
| assistantTextBuilder.insert(0, text); | |
| } | |
| } | |
| currentRole = role; | |
| } | |
| // Save any remaining pair | |
| String userText = userTextBuilder.toString().trim(); | |
| String assistantText = assistantTextBuilder.toString().trim(); | |
| if (!userText.isEmpty() && !assistantText.isEmpty()) { | |
| ConversationIndexMessage msg = ConversationIndexMessage | |
| .conversationIndexMessageBuilder() | |
| .type(appType) | |
| .question(userText) | |
| .response(assistantText) | |
| .finalAnswer(true) | |
| .sessionId(sessionId) | |
| .build(); | |
| messagePairs.add(msg); | |
| } | |
| // Reverse to maintain chronological order | |
| Collections.reverse(messagePairs); | |
| return messagePairs; | |
| } |
🤖 Prompt for AI Agents
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java`
around lines 1319 - 1417, The method extractMessagePairs currently assumes
messages is non-null and will NPE; add an early guard at the start of
extractMessagePairs to return Collections.emptyList() (or new ArrayList<>())
when messages is null or messages.isEmpty(), ensuring all callers receive an
empty list instead of throwing; update the beginning of the
extractMessagePairs(List<Message> messages, String sessionId, String appType)
method accordingly.
| // Check if agent uses unified interface (has model field) | ||
| // Unified interface agents store and retrieve structured messages with function calling | ||
| boolean usesUnifiedInterface = mlAgent.getModel() != null; | ||
|
|
||
| if (usesUnifiedInterface) { | ||
| // Get history first, then save new input messages | ||
| memory.getStructuredMessages(ActionListener.wrap(allMessages -> { | ||
| // Apply history limit | ||
| List<Message> history = messageHistoryLimit > 0 && allMessages.size() > messageHistoryLimit | ||
| ? allMessages.subList(allMessages.size() - messageHistoryLimit, allMessages.size()) | ||
| : allMessages; | ||
|
|
||
| // Save input messages | ||
| memory.saveStructuredMessages(inputMessages, ActionListener.wrap(v -> { | ||
| if (!history.isEmpty()) { | ||
| // Format history messages using the model provider for API-compatible output | ||
| ModelProvider modelProvider = ModelProviderFactory.getProvider(mlAgent.getModel().getModelProvider()); | ||
| MLAgentType agentType = MLAgentType.from(mlAgent.getType()); | ||
| Map<String, String> historyParams = modelProvider.mapMessages(history, agentType); | ||
| String formattedHistory = historyParams.get("body"); | ||
| if (formattedHistory != null && !formattedHistory.isEmpty()) { | ||
| params.put(NEW_CHAT_HISTORY, formattedHistory + ", "); | ||
| } | ||
| } | ||
| runAgent(mlAgent, params, listener, memory, functionCalling); | ||
| }, e -> { | ||
| log.error("Failed to save input messages", e); | ||
| listener.onFailure(e); | ||
| })); | ||
| }, e -> { | ||
| log.error("Failed to get history", e); | ||
| listener.onFailure(e); | ||
| })); |
There was a problem hiding this comment.
Validate unified-interface inputs before memory operations.
This path assumes inputMessages is non-null and that a valid ModelProvider exists; otherwise it can throw or fail with unclear errors. A fast failure makes misconfiguration explicit.
✅ Suggested guard for unified-interface inputs
- if (usesUnifiedInterface) {
+ if (usesUnifiedInterface) {
+ if (inputMessages == null || inputMessages.isEmpty()) {
+ listener.onFailure(new IllegalArgumentException("inputMessages is required for unified-interface agents"));
+ return;
+ }
+ ModelProvider modelProvider = ModelProviderFactory.getProvider(mlAgent.getModel().getModelProvider());
+ if (modelProvider == null) {
+ listener.onFailure(
+ new IllegalArgumentException("Unknown model provider: " + mlAgent.getModel().getModelProvider())
+ );
+ return;
+ }
// Get history first, then save new input messages
memory.getStructuredMessages(ActionListener.wrap(allMessages -> {
// Apply history limit
List<Message> history = messageHistoryLimit > 0 && allMessages.size() > messageHistoryLimit
? allMessages.subList(allMessages.size() - messageHistoryLimit, allMessages.size())
: allMessages;
// Save input messages
memory.saveStructuredMessages(inputMessages, ActionListener.wrap(v -> {
if (!history.isEmpty()) {
// Format history messages using the model provider for API-compatible output
- ModelProvider modelProvider = ModelProviderFactory.getProvider(mlAgent.getModel().getModelProvider());
MLAgentType agentType = MLAgentType.from(mlAgent.getType());
Map<String, String> historyParams = modelProvider.mapMessages(history, agentType);🤖 Prompt for AI Agents
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java`
around lines 271 - 303, The unified-interface branch (usesUnifiedInterface in
MLChatAgentRunner) assumes inputMessages is non-null and that
ModelProviderFactory.getProvider(mlAgent.getModel().getModelProvider()) returns
a valid provider; add quick validation before calling
memory.getStructuredMessages: check inputMessages != null and not empty, verify
mlAgent.getModel() and mlAgent.getModel().getModelProvider() are present, call
ModelProviderFactory.getProvider(...) and fail-fast with listener.onFailure(...)
if it returns null or throws; only then proceed to
memory.getStructuredMessages(...) / memory.saveStructuredMessages(...) and
runAgent(...).
| public void saveStructuredMessages(List<Message> messages, ActionListener<Void> listener) { | ||
| log | ||
| .debug( | ||
| "saveStructuredMessages: Entry - memoryContainerId={}, conversationId={}, messages count={}", | ||
| memoryContainerId, | ||
| conversationId, | ||
| messages != null ? messages.size() : "null" | ||
| ); | ||
| if (Strings.isNullOrEmpty(memoryContainerId)) { | ||
| listener.onFailure(new IllegalStateException("Memory container ID is not configured for this AgenticConversationMemory")); | ||
| return; | ||
| } | ||
|
|
||
| if (messages == null || messages.isEmpty()) { | ||
| listener.onResponse(null); | ||
| return; | ||
| } | ||
|
|
||
| AtomicInteger remaining = new AtomicInteger(messages.size()); | ||
| AtomicBoolean hasError = new AtomicBoolean(false); | ||
|
|
||
| for (int i = 0; i < messages.size(); i++) { | ||
| Message message = messages.get(i); | ||
|
|
||
| // Build namespace | ||
| Map<String, String> namespace = new HashMap<>(); | ||
| namespace.put(SESSION_ID_FIELD, conversationId); | ||
|
|
||
| // Build structured_data_blob | ||
| Map<String, Object> structuredData = new HashMap<>(); | ||
| Map<String, Object> serializableMessage = gson.fromJson(StringUtils.toJson(message), new TypeToken<Map<String, Object>>() { | ||
| }.getType()); | ||
| structuredData.put("message", serializableMessage); | ||
|
|
||
| // Build metadata | ||
| Map<String, String> metadata = new HashMap<>(); | ||
| metadata.put("type", "structured_message"); | ||
| if (message.getRole() != null) { | ||
| metadata.put("role", message.getRole()); | ||
| } | ||
|
|
||
| // Use messageId as sequence number so retrieval can sort by | ||
| // (created_time ASC, message_id ASC) to preserve ordering. | ||
| MLAddMemoriesInput input = MLAddMemoriesInput | ||
| .builder() | ||
| .memoryContainerId(memoryContainerId) | ||
| .structuredDataBlob(structuredData) | ||
| .messageId(i) | ||
| .namespace(namespace) | ||
| .metadata(metadata) | ||
| .infer(false) | ||
| .build(); | ||
|
|
||
| MLAddMemoriesRequest request = MLAddMemoriesRequest.builder().mlAddMemoryInput(input).build(); | ||
|
|
||
| int index = i; | ||
| client.execute(MLAddMemoriesAction.INSTANCE, request, ActionListener.wrap(response -> { | ||
| log.debug("Saved structured message {} of {} to session {}", index + 1, messages.size(), conversationId); | ||
| if (remaining.decrementAndGet() == 0) { | ||
| if (hasError.get()) { | ||
| listener.onFailure(new RuntimeException("One or more structured messages failed to save")); | ||
| } else { | ||
| listener.onResponse(null); | ||
| } | ||
| } | ||
| }, e -> { | ||
| log.error("Failed to save structured message {} of {} to session {}", index + 1, messages.size(), conversationId, e); | ||
| hasError.set(true); | ||
| if (remaining.decrementAndGet() == 0) { | ||
| listener.onFailure(e); | ||
| } | ||
| })); | ||
| } | ||
| } |
There was a problem hiding this comment.
Preserve message ordering when saving structured messages.
Writes are issued concurrently, but retrieval sorts by created_time then message_id. If requests arrive out of order (or different batches share the same timestamp), chat history can be reordered. Consider serializing saves (or persisting a monotonic sequence) to keep ordering stable.
✅ Suggested sequential save to preserve ordering
- AtomicInteger remaining = new AtomicInteger(messages.size());
- AtomicBoolean hasError = new AtomicBoolean(false);
-
- for (int i = 0; i < messages.size(); i++) {
- Message message = messages.get(i);
- ...
- client.execute(MLAddMemoriesAction.INSTANCE, request, ActionListener.wrap(response -> {
- log.debug("Saved structured message {} of {} to session {}", index + 1, messages.size(), conversationId);
- if (remaining.decrementAndGet() == 0) {
- if (hasError.get()) {
- listener.onFailure(new RuntimeException("One or more structured messages failed to save"));
- } else {
- listener.onResponse(null);
- }
- }
- }, e -> {
- log.error("Failed to save structured message {} of {} to session {}", index + 1, messages.size(), conversationId, e);
- hasError.set(true);
- if (remaining.decrementAndGet() == 0) {
- listener.onFailure(e);
- }
- }));
- }
+ saveStructuredMessagesSequentially(messages, 0, new AtomicBoolean(false), listener);+ private void saveStructuredMessagesSequentially(
+ List<Message> messages,
+ int index,
+ AtomicBoolean hasError,
+ ActionListener<Void> listener
+ ) {
+ if (index >= messages.size()) {
+ if (hasError.get()) {
+ listener.onFailure(new RuntimeException("One or more structured messages failed to save"));
+ } else {
+ listener.onResponse(null);
+ }
+ return;
+ }
+
+ Message message = messages.get(index);
+ // build namespace/structuredData/metadata/request as before...
+ client.execute(MLAddMemoriesAction.INSTANCE, request, ActionListener.wrap(response -> {
+ log.debug("Saved structured message {} of {} to session {}", index + 1, messages.size(), conversationId);
+ saveStructuredMessagesSequentially(messages, index + 1, hasError, listener);
+ }, e -> {
+ log.error("Failed to save structured message {} of {} to session {}", index + 1, messages.size(), conversationId, e);
+ hasError.set(true);
+ saveStructuredMessagesSequentially(messages, index + 1, hasError, listener);
+ }));
+ }🤖 Prompt for AI Agents
In
`@ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/AgenticConversationMemory.java`
around lines 429 - 502, saveStructuredMessages issues concurrent writes which
can arrive out of order despite using messageId=i; to fix, make saves sequential
or assign a strictly monotonic sequence number instead of relying on loop index.
Modify saveStructuredMessages to either: 1) replace the parallel client.execute
loop with a recursive/iterative sequential publisher (e.g., a private helper
saveNext(int index) that builds MLAddMemoriesInput/MLAddMemoriesRequest and
calls client.execute and on success calls saveNext(index+1), propagating
failures to listener), or 2) keep parallel writes but generate and persist a
monotonic sequence number for MLAddMemoriesInput.messageId (e.g., from an
AtomicLong sequence obtained from the memory container) so ordering is
guaranteed; adjust handling around MLAddMemoriesInput, MLAddMemoriesRequest,
client.execute, remaining/hasError, and messageId accordingly.
|
Can you add sample request/response |
|
Can you please update the PR description with more details? |
|
@ylwu-amzn @dhrubo-os added context and sample request/response in description, could you please help take a look again? This PR also addresses bug #4595 |
|
From the test example, seem |
Using a combination of created time and message id to sort the messages. Reason is that when using just But with this combination, we do not need a get request to store new messages as newer messages have higher timestamps. Messages stored at the same time will have different message IDs. You can see different IDs in the tool result sample response. |
jiapingzeng
force-pushed
the
extend-memory-interface
branch
from
4d9f496 to
7c1e93d
Compare
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
… not exist Signed-off-by: Jiaping Zeng <jpz@amazon.com>
…cutor Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
…AgentExecutor" This reverts commit a364488. Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
…nput" This reverts commit 1e667ee. Signed-off-by: Jiaping Zeng <jpz@amazon.com>
jiapingzeng
force-pushed
the
extend-memory-interface
branch
from
004eadc to
12c4fc6
Compare
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
jiapingzeng
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #4645 +/- ##
============================================
- Coverage 77.23% 77.06% -0.18%
- Complexity 11241 11368 +127
============================================
Files 944 944
Lines 50459 51144 +685
Branches 6073 6197 +124
============================================
+ Hits 38973 39415 +442
- Misses 8930 9118 +188
- Partials 2556 2611 +55
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:
|
Signed-off-by: Jiaping Zeng <jpz@amazon.com>
PR Reviewer Guide 🔍Here are some key observations to aid the review process:
|
PR Code Suggestions ✨Explore these optional code suggestions:
|
|
With the latest changes, all memory types are supported for unified interface agent. However, trace is disabled for unified interface agent and |
pyek-bot
left a comment
pyek-bot
left a comment
There was a problem hiding this comment.
Approving PR to unblock. Most changes look good. I will handle the gap when building the v2 changes for unified interface.
5 participants
Description
As a part of #4552, we introduced support for messages array as input for agent execute. However, in the initial implementation, we only stored plain text in memory. Multimodal messages and tool use/results are dropped. This is inline with our current implementation of the chat agent where a new request would only have access to initial user question and final assistant response from the previous request.
With this PR, we will also store multimodal messages and tool results from previous requests so that new requests will have full context. In addition, this PR moves memory storing logic from MLAgentExecutor to the memory layer by introducing new methods in the memory interface that handles messages array, so that each memory type can handle memory according to its capabilities. i.e. ConversationIndexMemory will still only store text messages whereas AgenticConversationMemory and RemoteAgenticMemory will store all messages.
Sample request/response:
Agent here is a conversational agent with agentic memory.
Checking memory:
Follow-up question:
Check memory again:
With tool result: