refactor(memory): manage state serialization via StateCenter in MemoryUnit, and support optional loading on register in StateCenter

Partner-Core/src/main/java/work/slhaf/partner/core/memory/MemoryCore.java

@@ -45,7 +45,13 @@ public class MemoryCore implements StateSerializable {
 
     @CapabilityMethod
     public MemoryUnit getMemoryUnit(String unitId) {
-        return memoryUnits.computeIfAbsent(unitId, MemoryUnit::new);
+        MemoryUnit unit = memoryUnits.computeIfAbsent(unitId, id -> {
+            MemoryUnit newUnit = new MemoryUnit(id);
+            newUnit.register();
+            return newUnit;
+        });
+        unit.load();
+        return unit;
     }
 
     @CapabilityMethod

Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemorySlice.java

@@ -1,18 +1,11 @@
 package work.slhaf.partner.core.memory.pojo;
 
-import lombok.Data;
-import lombok.EqualsAndHashCode;
-import work.slhaf.partner.framework.agent.common.entity.PersistableObject;
+import lombok.Getter;
 
-import java.io.Serial;
 import java.util.UUID;
 
-@EqualsAndHashCode(callSuper = true)
-@Data
-public class MemorySlice extends PersistableObject implements Comparable<MemorySlice> {
-
-    @Serial
-    private static final long serialVersionUID = 1L;
+@Getter
+public class MemorySlice implements Comparable<MemorySlice> {
 
     private final String id;
     private final Integer startIndex;
@@ -28,6 +21,18 @@ public class MemorySlice extends PersistableObject implements Comparable<MemoryS
         this.summary = summary;
     }
 
+    private MemorySlice(String id, Integer startIndex, Integer endIndex, String summary, Long timestamp) {
+        this.id = id;
+        this.startIndex = startIndex;
+        this.endIndex = endIndex;
+        this.summary = summary;
+        this.timestamp = timestamp;
+    }
+
+    public static MemorySlice restore(String id, Integer startIndex, Integer endIndex, String summary, Long timestamp) {
+        return new MemorySlice(id, startIndex, endIndex, summary, timestamp);
+    }
+
     @Override
     public int compareTo(MemorySlice memorySlice) {
         if (memorySlice.getTimestamp() > this.getTimestamp()) {

Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemoryUnit.java

@@ -1,24 +1,125 @@
 package work.slhaf.partner.core.memory.pojo;
 
+import com.alibaba.fastjson2.JSONArray;
+import com.alibaba.fastjson2.JSONObject;
 import lombok.Getter;
+import org.jetbrains.annotations.NotNull;
 import work.slhaf.partner.framework.agent.model.pojo.Message;
+import work.slhaf.partner.framework.agent.state.State;
+import work.slhaf.partner.framework.agent.state.StateSerializable;
+import work.slhaf.partner.framework.agent.state.StateValue;
 
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 @Getter
-public class MemoryUnit {
+public class MemoryUnit implements StateSerializable {
 
     private final String id;
     private final List<Message> conversationMessages = new ArrayList<>();
-    private Long timestamp;
+    private Long timestamp = 0L;
     private final List<MemorySlice> slices = new ArrayList<>();
 
     public MemoryUnit(String id) {
         this.id = id;
+        this.register();
     }
 
     public void updateTimestamp() {
         timestamp = System.currentTimeMillis();
     }
+
+    @Override
+    public @NotNull Path statePath() {
+        return Path.of("core", "memory", "memory-unit" + id + ".json");
+    }
+
+    @Override
+    public void load(@NotNull JSONObject state) {
+        Long loadedTimestamp = state.getLong("update_timestamp");
+        this.timestamp = loadedTimestamp != null ? loadedTimestamp : 0L;
+
+        this.conversationMessages.clear();
+        this.slices.clear();
+
+        JSONArray messageArray = state.getJSONArray("conversation_messages");
+        if (messageArray != null) {
+            for (int i = 0; i < messageArray.size(); i++) {
+                JSONObject messageObject = messageArray.getJSONObject(i);
+                if (messageObject == null) {
+                    continue;
+                }
+
+                String role = messageObject.getString("role");
+                String content = messageObject.getString("content");
+                if (role == null || content == null) {
+                    continue;
+                }
+
+                Message message = new Message(Message.Character.fromValue(role), content);
+
+                this.conversationMessages.add(message);
+            }
+        }
+
+        var sliceArray = state.getJSONArray("memory_slices");
+        if (sliceArray != null) {
+            for (int i = 0; i < sliceArray.size(); i++) {
+                JSONObject sliceObject = sliceArray.getJSONObject(i);
+                if (sliceObject == null) {
+                    continue;
+                }
+
+                String sliceId = sliceObject.getString("id");
+                Integer startIndex = sliceObject.getInteger("start_index");
+                Integer endIndex = sliceObject.getInteger("end_index");
+                String summary = sliceObject.getString("summary");
+                Long createdTimestamp = sliceObject.getLong("created_timestamp");
+
+                if (sliceId == null || startIndex == null || endIndex == null || summary == null || createdTimestamp == null) {
+                    continue;
+                }
+
+                MemorySlice slice = MemorySlice.restore(sliceId, startIndex, endIndex, summary, createdTimestamp);
+
+                this.slices.add(slice);
+            }
+        }
+    }
+
+    @Override
+    public @NotNull State convert() {
+        State state = new State();
+        state.append("id", StateValue.str(id));
+        state.append("update_timestamp", StateValue.num(timestamp));
+
+        List<StateValue.Obj> convertedMessageList = conversationMessages.stream().map(message -> {
+            Map<String, StateValue> convertedMap = Map.of(
+                    "role", StateValue.str(message.roleValue()),
+                    "content", StateValue.str(message.getContent())
+            );
+            return StateValue.obj(convertedMap);
+        }).toList();
+        state.append("conversation_messages", StateValue.arr(convertedMessageList));
+
+        List<StateValue.Obj> convertedSliceList = slices.stream().map(slice -> {
+            Map<String, StateValue> convertedMap = Map.of(
+                    "id", StateValue.str(slice.getId()),
+                    "start_index", StateValue.num(slice.getStartIndex()),
+                    "end_index", StateValue.num(slice.getEndIndex()),
+                    "summary", StateValue.str(slice.getSummary()),
+                    "created_timestamp", StateValue.num(slice.getTimestamp())
+            );
+            return StateValue.obj(convertedMap);
+        }).toList();
+        state.append("memory_slices", StateValue.arr(convertedSliceList));
+        return state;
+    }
+
+    @Override
+    public boolean autoLoadOnRegister() {
+        return false;
+    }
 }