From 2458ea48492f6007b9884d508ed88b80eb9016fb Mon Sep 17 00:00:00 2001 From: slhafzjw Date: Tue, 7 Apr 2026 12:27:36 +0800 Subject: [PATCH] refactor(memory): manage state serialization via StateCenter in MemoryUnit, and support optional loading on register in StateCenter --- .../slhaf/partner/core/memory/MemoryCore.java | 8 +- .../partner/core/memory/pojo/MemorySlice.java | 25 +++-- .../partner/core/memory/pojo/MemoryUnit.java | 105 +++++++++++++++++- .../framework/agent/state/StateCenter.kt | 73 ++++++++++-- 4 files changed, 188 insertions(+), 23 deletions(-) diff --git a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/MemoryCore.java b/Partner-Core/src/main/java/work/slhaf/partner/core/memory/MemoryCore.java index 53ac9feb..5b551ebb 100644 --- a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/MemoryCore.java +++ b/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 diff --git a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemorySlice.java b/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemorySlice.java index f46d7e7c..5cdc1cda 100644 --- a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemorySlice.java +++ b/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 { - - @Serial - private static final long serialVersionUID = 1L; +@Getter +public class MemorySlice implements Comparable { private final String id; private final Integer startIndex; @@ -28,6 +21,18 @@ public class MemorySlice extends PersistableObject implements Comparable this.getTimestamp()) { diff --git a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemoryUnit.java b/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemoryUnit.java index 383c9c1a..06b68fc1 100644 --- a/Partner-Core/src/main/java/work/slhaf/partner/core/memory/pojo/MemoryUnit.java +++ b/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 conversationMessages = new ArrayList<>(); - private Long timestamp; + private Long timestamp = 0L; private final List 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 convertedMessageList = conversationMessages.stream().map(message -> { + Map 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 convertedSliceList = slices.stream().map(slice -> { + Map 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; + } } diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/state/StateCenter.kt b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/state/StateCenter.kt index 74fee479..1d08694a 100644 --- a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/state/StateCenter.kt +++ b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/state/StateCenter.kt @@ -2,6 +2,7 @@ package work.slhaf.partner.framework.agent.state import com.alibaba.fastjson2.JSONArray import com.alibaba.fastjson2.JSONObject +import org.slf4j.LoggerFactory import work.slhaf.partner.framework.agent.config.ConfigCenter import java.nio.file.Files import java.nio.file.Path @@ -13,7 +14,9 @@ import kotlin.io.path.writeText object StateCenter { - private val stateRegistry = ConcurrentHashMap() + private val log = LoggerFactory.getLogger(StateCenter::class.java) + + private val stateRegistry = ConcurrentHashMap() fun register(stateSerializable: StateSerializable): JSONObject? { val relativePath = stateSerializable.statePath().normalize() @@ -23,8 +26,9 @@ object StateCenter { val finalStatePath = stateDir.resolve(relativePath).normalize() check(finalStatePath.startsWith(stateDir)) { "StatePath escapes stateDir" } - val previous = stateRegistry.putIfAbsent(finalStatePath, stateSerializable) - check(previous == null || previous === stateSerializable) { + val stateRecord = StateRecord(stateSerializable) + val previous = stateRegistry.putIfAbsent(finalStatePath, stateRecord) + check(previous == null || previous.serializable === stateSerializable) { "StatePath already registered: $finalStatePath" } @@ -35,15 +39,43 @@ object StateCenter { check(finalStatePath.isRegularFile()) { "StatePath must point to a regular file: $finalStatePath" } check(finalStatePath.toFile().canRead()) { "StateFile must be readable: $finalStatePath" } + if (!stateSerializable.autoLoadOnRegister()) { + return null + } + + stateRecord.loaded = true + return JSONObject.parseObject(finalStatePath.readText()) } - fun save() { - stateRegistry.forEach { (path, state) -> - path.parent?.let(Files::createDirectories) - path.writeText(state.convert().toString()) + fun load(path: Path) { + val finalStatePath = ConfigCenter.paths.stateDir.normalize().resolve(path).normalize() + if (!stateRegistry.containsKey(path)) { + return + } + val record = stateRegistry[finalStatePath] ?: return + record.loaded = true + if (!finalStatePath.exists()) { + return + } + try { + val json = JSONObject.parseObject(finalStatePath.readText()) + record.serializable.load(json) + } catch (_: Exception) { + log.warn("StateCenter loading failed: $path") } } + + fun save() { + stateRegistry.forEach { (path, record) -> + if (!record.loaded) { + return@forEach + } + path.parent?.let(Files::createDirectories) + path.writeText(record.serializable.convert().toString()) + } + } + } interface StateSerializable { @@ -57,9 +89,27 @@ interface StateSerializable { fun statePath(): Path + /** + * 手动加载状态数据 + */ + fun load() { + StateCenter.load(statePath()) + } + + /** + * 状态加载逻辑 + */ fun load(state: JSONObject) + /** + * 数据转换为状态逻辑 + */ fun convert(): State + + /** + * 是否在注册时即触发一次加载 + */ + fun autoLoadOnRegister(): Boolean = true } class State { @@ -68,8 +118,6 @@ class State { fun append(key: String, value: StateValue) = json.put(key, value.toJsonValue()) - fun toJson(): JSONObject = json - override fun toString(): String = json.toString() private fun StateValue.toJsonValue(): Any = when (this) { @@ -104,4 +152,9 @@ sealed interface StateValue { @JvmStatic fun obj(value: Map) = Obj(value) } -} \ No newline at end of file +} + +data class StateRecord( + val serializable: StateSerializable, + var loaded: Boolean = false +) \ No newline at end of file