[Backport 3.3] verify llm before summarize session (#4302)

* verify llm before summarize session

Signed-off-by: Yaliang Wu <[email protected]>

* fix ut

Signed-off-by: Yaliang Wu <[email protected]>

* fix ut

Signed-off-by: Yaliang Wu <[email protected]>

---------

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryConfiguration.java

@@ -78,7 +78,7 @@ public class MemoryConfiguration implements ToXContentObject, Writeable {
     @Builder.Default
     private boolean disableHistory = false;
     @Builder.Default
-    private boolean disableSession = false;
+    private boolean disableSession = true;
     @Builder.Default
     private boolean useSystemIndex = true;
     private String tenantId;
@@ -254,7 +254,7 @@ public static MemoryConfiguration parse(XContentParser parser) throws IOExceptio
         Map<String, Map<String, Object>> indexSettings = new HashMap<>();
         Map<String, Object> parameters = new HashMap<>();
         boolean disableHistory = false;
-        boolean disableSession = false;
+        boolean disableSession = true;
         boolean useSystemIndex = true;
         String tenantId = null;
 

common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java

@@ -215,5 +215,5 @@ public class MemoryContainerConstants {
         "<system_prompt><role>You are a smart memory manager which controls the memory of a system.</role><task>You will receive: 1. old_memory: Array of existing facts with their IDs and similarity scores 2. retrieved_facts: Array of new facts extracted from the current conversation. Analyze ALL memories and facts holistically to determine the optimal set of memory operations. Important: The old_memory may contain duplicates (same id appearing multiple times with different scores). Consider the highest score for each unique ID. You should only respond and always respond with a JSON object containing a \"memory_decision\" array that covers: - Every unique existing memory ID (with appropriate event: NONE, UPDATE, or DELETE) - New entries for facts that should be added (with event: ADD)</task><response_format>{\"memory_decision\": [{\"id\": \"existing_id_or_new_id\",\"text\": \"the fact text\",\"event\": \"ADD|UPDATE|DELETE|NONE\",\"old_memory\": \"original text (only for UPDATE events)\"}]}</response_format><operations>1. **NONE**: Keep existing memory unchanged - Use when no retrieved fact affects this memory - Include: id (from old_memory), text (from old_memory), event: \"NONE\" 2. **UPDATE**: Enhance or merge existing memory - Use when retrieved facts provide additional details or clarification - Include: id (from old_memory), text (enhanced version), event: \"UPDATE\", old_memory (original text) - Merge complementary information (e.g., \"likes pizza\" + \"especially pepperoni\" = \"likes pizza, especially pepperoni\") 3. **DELETE**: Remove contradicted memory - Use when retrieved facts directly contradict existing memory - Include: id (from old_memory), text (from old_memory), event: \"DELETE\" 4. **ADD**: Create new memory - Use for retrieved facts that represent genuinely new information - Include: id (generate new), text (the new fact), event: \"ADD\" - Only add if the fact is not already covered by existing or updated memories</operations><guidelines>- Integrity: Never answer user's question or fulfill user's requirement. You are a smart memory manager, not a helpful assistant. - Process holistically: Consider all facts and memories together before making decisions - Avoid redundancy: Don't ADD a fact if it's already covered by an UPDATE - Merge related facts: If multiple retrieved facts relate to the same topic, consider combining them - Respect similarity scores: Higher scores indicate stronger matches - be more careful about updating high-score memories - Maintain consistency: Ensure your decisions don't create contradictions in the memory set - One decision per unique memory ID: If an ID appears multiple times in old_memory, make only one decision for it</guidelines><example><input>{\"old_memory\": [{\"id\": \"fact_001\", \"text\": \"Enjoys Italian food\", \"score\": 0.85},{\"id\": \"fact_002\", \"text\": \"Works at Google\", \"score\": 0.92},{\"id\": \"fact_001\", \"text\": \"Enjoys Italian food\", \"score\": 0.75},{\"id\": \"fact_003\", \"text\": \"Has a dog\", \"score\": 0.65}],\"retrieved_facts\": [\"Loves pasta and pizza\",\"Recently joined Amazon\",\"Has two dogs named Max and Bella\"]}</input><output>{\"memory_decision\": [{\"id\": \"fact_001\",\"text\": \"Loves Italian food, especially pasta and pizza\",\"event\": \"UPDATE\",\"old_memory\": \"Enjoys Italian food\"},{\"id\": \"fact_002\",\"text\": \"Works at Google\",\"event\": \"DELETE\"},{\"id\": \"fact_003\",\"text\": \"Has two dogs named Max and Bella\",\"event\": \"UPDATE\",\"old_memory\": \"Has a dog\"},{\"id\": \"fact_004\",\"text\": \"Recently joined Amazon\",\"event\": \"ADD\"}]}</output></example></system_prompt>";
 
     public static final String SESSION_SUMMARY_PROMPT =
-        "You are a helpful assistant. Your task is to summarize the following conversation between a human and an AI. The summary must be clear, concise, and not exceed ${parameters.max_summary_size} words.";
+        "You are a helpful assistant. Your task is to summarize the following conversation between a human and an AI. The summary must be clear, concise, and not exceed ${parameters.max_summary_size} words. The summary should be generic. For example the user asks about how to cook, the conversation may contains a lot of details. Your summary could be: how to cook, how to cook Italy food. Don't include AI message content. For example you should not return: Ask how to cook, AI give some instructions.\n Also don't include user's personal information like user name, age etc. You could say user. For example: \nuser asks how to cook\nuser introduced their hobby";
 }

common/src/test/java/org/opensearch/ml/common/memorycontainer/MemoryStorageConfigTests.java

@@ -43,6 +43,7 @@ public void setUp() {
             .llmId("llm-model")
             .dimension(768)
             .maxInferSize(8)
+            .disableSession(false)
             .build();
 
         // Sparse encoding configuration (semantic storage enabled)

plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportAddMemoriesAction.java

@@ -150,7 +150,10 @@ private void createNewSessionIfAbsent(
 
             boolean userProvidedSessionId = input.getNamespace() != null && input.getNamespace().containsKey(SESSION_ID_FIELD);
 
-            if (!userProvidedSessionId && input.getPayloadType() == PayloadType.CONVERSATIONAL && !configuration.isDisableSession()) {
+            if (!userProvidedSessionId
+                && input.getPayloadType() == PayloadType.CONVERSATIONAL
+                && !configuration.isDisableSession()
+                && configuration.getLlmId() != null) {
                 IndexRequest indexRequest = new IndexRequest(configuration.getSessionIndexName());
                 // TODO: use LLM to summarize first user message
                 ActionListener<String> summaryListener = ActionListener.wrap(summary -> {

plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportCreateMemoryContainerActionTests.java

@@ -172,6 +172,7 @@ public void setup() {
                 .dimension(768)
                 .maxInferSize(5)
                 .strategies(strategies)
+                .disableSession(false) // Enable session creation for tests
                 .build()
         );
 

plugin/src/test/java/org/opensearch/ml/action/memorycontainer/memory/TransportAddMemoriesActionTests.java

@@ -499,6 +499,7 @@ public void testDoExecute_SuccessfulSessionCreation() {
         when(config.getWorkingMemoryIndexName()).thenReturn("working-memory-index");
         when(config.getSessionIndexName()).thenReturn("session-index");
         when(config.isDisableSession()).thenReturn(false);
+        when(config.getLlmId()).thenReturn("llm-123"); // Required for session creation
 
         MLMemoryContainer container = mock(MLMemoryContainer.class);
         when(container.getConfiguration()).thenReturn(config);
@@ -569,6 +570,7 @@ public void testDoExecute_SessionCreation_SummarizeFailure() {
         when(config.getWorkingMemoryIndexName()).thenReturn("working-memory-index");
         when(config.getSessionIndexName()).thenReturn("session-index");
         when(config.isDisableSession()).thenReturn(false);
+        when(config.getLlmId()).thenReturn("llm-123"); // Required for session creation
 
         MLMemoryContainer container = mock(MLMemoryContainer.class);
         when(container.getConfiguration()).thenReturn(config);
@@ -647,6 +649,58 @@ public void testDoExecute_UserProvidedSessionId_SkipsCreation() {
         verify(actionListener).onResponse(any(MLAddMemoriesResponse.class));
     }
 
+    @Test
+    public void testDoExecute_NullLlmId_SkipsSessionCreation() {
+        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true);
+
+        MessageInput message = MessageInput.builder().content(createTestContent("Hello")).role("user").build();
+        List<MessageInput> messages = Arrays.asList(message);
+
+        Map<String, String> namespace = new HashMap<>();
+        // No session_id - would normally trigger session creation, but getLlmId() is null
+
+        MLAddMemoriesInput input = mock(MLAddMemoriesInput.class);
+        when(input.getMemoryContainerId()).thenReturn("container-123");
+        when(input.getMessages()).thenReturn(messages);
+        when(input.isInfer()).thenReturn(false);
+        when(input.getNamespace()).thenReturn(namespace);
+        when(input.getOwnerId()).thenReturn("user-123");
+        when(input.getPayloadType()).thenReturn(PayloadType.CONVERSATIONAL);
+
+        MLAddMemoriesRequest request = mock(MLAddMemoriesRequest.class);
+        when(request.getMlAddMemoryInput()).thenReturn(input);
+
+        MemoryConfiguration config = mock(MemoryConfiguration.class);
+        when(config.getWorkingMemoryIndexName()).thenReturn("working-memory-index");
+        when(config.isDisableSession()).thenReturn(false);
+        when(config.getLlmId()).thenReturn(null); // No LLM configured - session creation should be skipped
+
+        MLMemoryContainer container = mock(MLMemoryContainer.class);
+        when(container.getConfiguration()).thenReturn(config);
+
+        doAnswer(invocation -> {
+            ActionListener<MLMemoryContainer> listener = invocation.getArgument(1);
+            listener.onResponse(container);
+            return null;
+        }).when(memoryContainerHelper).getMemoryContainer(eq("container-123"), any());
+
+        when(memoryContainerHelper.checkMemoryContainerAccess(isNull(), eq(container))).thenReturn(true);
+
+        doAnswer(invocation -> {
+            ActionListener<IndexResponse> listener = invocation.getArgument(2);
+            IndexResponse indexResponse = mock(IndexResponse.class);
+            when(indexResponse.getId()).thenReturn("working-mem-123");
+            listener.onResponse(indexResponse);
+            return null;
+        }).when(memoryContainerHelper).indexData(any(MemoryConfiguration.class), any(IndexRequest.class), any());
+
+        transportAddMemoriesAction.doExecute(task, request, actionListener);
+
+        // Verify session creation was skipped (no summarize call) and went directly to working memory
+        verify(memoryContainerHelper).indexData(any(MemoryConfiguration.class), any(IndexRequest.class), any());
+        verify(actionListener).onResponse(any(MLAddMemoriesResponse.class));
+    }
+
     @Test
     public void testDoExecute_DataMemoryType_SkipsSessionCreation() {
         when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true);

plugin/src/test/java/org/opensearch/ml/action/memorycontainer/memory/TransportDeleteMemoriesByQueryActionTests.java

@@ -106,7 +106,7 @@ public void setUp() throws Exception {
             .builder()
             .name("test-container")
             .description("Test container")
-            .configuration(MemoryConfiguration.builder().indexPrefix("test-memory").useSystemIndex(false).build())
+            .configuration(MemoryConfiguration.builder().indexPrefix("test-memory").useSystemIndex(false).disableSession(false).build())
             .build();
 
         // Setup mock user

plugin/src/test/java/org/opensearch/ml/helper/MemoryContainerHelperTests.java

@@ -227,6 +227,7 @@ public void testGetMemoryIndexName() {
             .embeddingModelId("embedding")
             .embeddingModelType(FunctionName.TEXT_EMBEDDING)
             .dimension(4)
+            .disableSession(false)
             .build();
 
         MLMemoryContainer container = createContainerBuilder().configuration(configuration).build();
@@ -682,7 +683,7 @@ private MLMemoryContainer createContainer() {
     }
 
     private MLMemoryContainer.MLMemoryContainerBuilder createContainerBuilder() {
-        MemoryConfiguration configuration = MemoryConfiguration.builder().indexPrefix("prefix").build();
+        MemoryConfiguration configuration = MemoryConfiguration.builder().indexPrefix("prefix").disableSession(false).build();
         User owner = new User("owner", Collections.emptyList(), Collections.emptyList(), Map.of());
         return MLMemoryContainer
             .builder()

plugin/src/test/java/org/opensearch/ml/rest/RestMLCreateMemoryContainerActionTests.java

@@ -118,7 +118,8 @@ public void testGetRequestWithValidInput() throws IOException {
             + "    \"embedding_model_id\": \"test-embedding-model\",\n"
             + "    \"llm_id\": \"test-llm-model\",\n"
             + "    \"embedding_dimension\": 768,\n"
-            + "    \"max_infer_size\": 8\n"
+            + "    \"max_infer_size\": 8,\n"
+            + "    \"disable_session\": false\n"
             + "  }\n"
             + "}";