fix(rag): Add missing payload for ElasticsearchStore

agentscope-extensions/agentscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/ElasticsearchStore.java

@@ -34,13 +34,15 @@
 import co.elastic.clients.transport.ElasticsearchTransport;
 import co.elastic.clients.transport.rest5_client.Rest5ClientTransport;
 import co.elastic.clients.transport.rest5_client.low_level.Rest5Client;
+import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.agentscope.core.message.ContentBlock;
 import io.agentscope.core.message.TextBlock;
 import io.agentscope.core.rag.exception.VectorStoreException;
 import io.agentscope.core.rag.model.Document;
 import io.agentscope.core.rag.model.DocumentMetadata;
 import io.agentscope.core.rag.store.dto.SearchDocumentDto;
+import io.agentscope.core.util.JsonUtils;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -93,6 +95,7 @@ public class ElasticsearchStore implements VDBStoreBase, AutoCloseable {
     private static final String FIELD_DOC_ID = "doc_id";
     private static final String FIELD_CHUNK_ID = "chunk_id";
     private static final String FIELD_CONTENT = "content";
+    private static final String FIELD_PAYLOAD = "payload";
 
     private final String indexName;
     private final int dimensions;
@@ -362,6 +365,7 @@ private void ensureIndex() throws VectorStoreException {
             properties.put(FIELD_CHUNK_ID, idProperty);
             properties.put(FIELD_CONTENT, contentProperty);
             properties.put(FIELD_VECTOR, vectorProperty);
+            properties.put(FIELD_PAYLOAD, contentProperty);
 
             TypeMapping mapping = new TypeMapping.Builder().properties(properties).build();
 
@@ -414,6 +418,11 @@ private Map<String, Object> mapToEsDocument(Document doc) {
             map.put(FIELD_CONTENT, meta.getContentText());
         }
 
+        Map<String, Object> customPayload = meta.getPayload();
+        if (customPayload != null && !customPayload.isEmpty()) {
+            map.put(FIELD_PAYLOAD, JsonUtils.getJsonCodec().toJson(customPayload));
+        }
+
         return map;
     }
 
@@ -426,6 +435,7 @@ private Document mapFromEsHit(Hit<Map> hit) {
             String docId = (String) source.get(FIELD_DOC_ID);
             String chunkId = (String) source.get(FIELD_CHUNK_ID);
             String contentJson = (String) source.get(FIELD_CONTENT);
+            String payloadJson = (String) source.get(FIELD_PAYLOAD);
 
             // Reconstruct ContentBlock
             ContentBlock content;
@@ -436,7 +446,19 @@ private Document mapFromEsHit(Hit<Map> hit) {
                 content = TextBlock.builder().text(contentJson).build();
             }
 
-            DocumentMetadata metadata = new DocumentMetadata(content, docId, chunkId);
+            Map<String, Object> customPayload = new HashMap<>();
+            if (payloadJson != null && !payloadJson.isBlank()) {
+                try {
+                    customPayload =
+                            JsonUtils.getJsonCodec()
+                                    .fromJson(payloadJson, new TypeReference<>() {});
+                } catch (Exception e) {
+                    log.warn("Failed to deserialize payload, using empty map", e);
+                }
+            }
+
+            DocumentMetadata metadata =
+                    new DocumentMetadata(content, docId, chunkId, customPayload);
             Document doc = new Document(metadata);
 
             // Set score if present

agentscope-extensions/agentscope-extensions-rag-simple/src/test/java/io/agentscope/core/rag/store/ElasticsearchStoreTest.java

@@ -376,6 +376,20 @@ void testAddDimensionMismatch() throws VectorStoreException {
                 .verify();
     }
 
+    @Test
+    @DisplayName("Should add documents with payload")
+    void testAddPayload() throws VectorStoreException {
+        store = createMockStoreForAdd(true);
+
+        TextBlock content = TextBlock.builder().text("Test content").build();
+        DocumentMetadata metadata =
+                new DocumentMetadata(content, "doc-1", "chunk-1", Map.of("k", "v"));
+        Document doc = new Document(metadata);
+        doc.setEmbedding(new double[TEST_DIMENSIONS]);
+
+        StepVerifier.create(store.add(List.of(doc))).verifyComplete();
+    }
+
     // ==================== Search Method Tests ====================
 
     @SuppressWarnings("unchecked")
@@ -433,6 +447,7 @@ void testSearchSuccess() throws VectorStoreException {
         List<Double> vector = new ArrayList<>();
         for (int i = 0; i < TEST_DIMENSIONS; i++) vector.add(0.0);
         source.put("vector", vector);
+        source.put("payload", "{\"k\":\"v\"}");
 
         store = createMockStoreForSearch(List.of(source));
 
@@ -450,6 +465,7 @@ void testSearchSuccess() throws VectorStoreException {
                             assertEquals(1, results.size());
                             assertEquals("doc-1", results.get(0).getMetadata().getDocId());
                             assertEquals(0.95, results.get(0).getScore());
+                            assertEquals("v", results.get(0).getMetadata().getPayloadValue("k"));
                         })
                 .verifyComplete();
     }