diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/beans/AgenticAggregateRuntimeFactory.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/beans/AgenticAggregateRuntimeFactory.java index 0dba4f4c..46a73823 100644 --- a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/beans/AgenticAggregateRuntimeFactory.java +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/beans/AgenticAggregateRuntimeFactory.java @@ -134,7 +134,8 @@ public AgenticAggregateRuntime getObject() { KafkaAggregateRuntime kafkaRuntime = createRuntime(agenticInfo, aggregate, agentPlatform); return new KafkaAgenticAggregateRuntime( - kafkaRuntime, objectMapper, agenticInfo.stateClass(), agentPlatform, aggregate); + kafkaRuntime, objectMapper, agenticInfo.stateClass(), agentPlatform, aggregate, + agenticInfo.maxTotalMemories(), agenticInfo.maxMemoriesAdded()); } @Override diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/DefaultAgent.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/DefaultAgent.java index 7d2c1618..f9d28569 100644 --- a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/DefaultAgent.java +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/DefaultAgent.java @@ -35,5 +35,5 @@ public class DefaultAgent { * The well-known agent name used to locate this default agent on the * {@link com.embabel.agent.core.AgentPlatform} during fallback resolution. */ - public static final String AGENT_NAME = "DefaultAgent"; + public static final String AGENT_NAME = "Default"; } diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationInput.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationInput.java new file mode 100644 index 00000000..d132c461 --- /dev/null +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationInput.java @@ -0,0 +1,48 @@ +/* + * Copyright 2022 - 2026 The Original Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.elasticsoftware.akces.agentic.embabel; + +import com.embabel.agent.core.ActionInvocation; +import org.elasticsoftware.akces.aggregate.AgenticAggregateMemory; +import org.elasticsoftware.akces.aggregate.AssignedTask; + +import java.util.List; + +/** + * Strongly-typed input for the {@link MemoryDistillerAgent} action method. + * + *

This record is placed on the Embabel Blackboard by the Akces runtime + * and injected into the agent's action method by the Embabel framework + * via type-based parameter resolution. + * + * @param agentTask the assigned task that triggered the completed process + * @param history the execution history (action invocations) from the completed process + * @param blackboardObjects all objects from the completed process's blackboard + * @param existingMemories the current list of memories from the aggregate state + * @param maxTotalMemories the total memory capacity of the aggregate + * @param maxMemoriesAdded the per-distillation budget for net new memories + */ +public record MemoryDistillationInput( + AssignedTask agentTask, + List history, + List blackboardObjects, + List existingMemories, + int maxTotalMemories, + int maxMemoriesAdded +) { +} diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResult.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResult.java new file mode 100644 index 00000000..873a7280 --- /dev/null +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResult.java @@ -0,0 +1,42 @@ +/* + * Copyright 2022 - 2026 The Original Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.elasticsoftware.akces.agentic.embabel; + +import org.elasticsoftware.akces.agentic.events.MemoryRevokedEvent; +import org.elasticsoftware.akces.agentic.events.MemoryStoredEvent; + +import java.util.List; + +/** + * Result object produced by the {@link MemoryDistillerAgent} after distilling + * relevant memories from a completed {@link com.embabel.agent.core.AgentProcess}. + * + *

Contains two lists: + *

    + *
  • {@code stored} — new memory entries as {@link MemoryStoredEvent}s
    • + *
  • {@code revoked} — memory entries to remove as {@link MemoryRevokedEvent}s
    • + *
+ * + * @param stored the list of memories to store + * @param revoked the list of memories to revoke + */ +public record MemoryDistillationResult( + List stored, + List revoked +) { +} diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillerAgent.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillerAgent.java new file mode 100644 index 00000000..bc759eb6 --- /dev/null +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillerAgent.java @@ -0,0 +1,197 @@ +/* + * Copyright 2022 - 2026 The Original Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.elasticsoftware.akces.agentic.embabel; + +import com.embabel.agent.api.annotation.AchievesGoal; +import com.embabel.agent.api.annotation.Action; +import com.embabel.agent.api.annotation.Agent; +import com.embabel.agent.api.common.ActionContext; +import com.embabel.agent.api.common.PlannerType; +import tools.jackson.databind.ObjectMapper; + +import java.util.List; +import java.util.Objects; + +/** + * Embabel agent responsible for distilling relevant memories from a successfully + * completed {@link com.embabel.agent.core.AgentProcess}. + * + *

When an agent process finishes with status {@code COMPLETED}, this agent is + * invoked with the process's execution history and all objects from its blackboard. + * It uses an LLM to analyze the process results and determine which memories should + * be stored and which existing memories should be revoked. + * + *

The agent produces a {@link MemoryDistillationResult} containing: + *

    + *
  • A list of new memories to store
    • + *
  • A list of existing memories to revoke
    • + *
+ * + *

The net number of new memories (stored minus revoked) is constrained by both + * the {@code maxTotalMemories} capacity and the per-distillation budget + * {@code maxMemoriesAdded}. + */ +@Agent(name = MemoryDistillerAgent.AGENT_NAME, + description = "Distills relevant memories from a completed agent process", + planner = PlannerType.GOAP) +public class MemoryDistillerAgent { + + /** + * The well-known agent name used to locate this agent on the + * {@link com.embabel.agent.core.AgentPlatform}. + */ + public static final String AGENT_NAME = "MemoryDistiller"; + + private final ObjectMapper objectMapper; + + /** + * Creates a new {@code MemoryDistillerAgent}. + * + * @param objectMapper the Jackson {@link ObjectMapper} used for JSON serialization + * of blackboard objects and memories in the LLM prompt + */ + public MemoryDistillerAgent(ObjectMapper objectMapper) { + this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null"); + } + + /** + * Distills memories from a completed agent process by analyzing the process + * history and blackboard contents using an LLM. + * + *

The Embabel framework injects the {@link MemoryDistillationInput} parameter + * from the Blackboard via type-based resolution. The input contains: + *

    + *
  • {@code agentTask} — the {@link org.elasticsoftware.akces.aggregate.AssignedTask} + * that triggered the completed process
    • + *
  • {@code history} — the execution history from the completed process
    • + *
  • {@code blackboardObjects} — all objects from the completed process's blackboard
    • + *
  • {@code existingMemories} — the current list of memories from the aggregate state
    • + *
  • {@code maxTotalMemories} — the total memory capacity of the aggregate
    • + *
  • {@code maxMemoriesAdded} — the per-distillation budget for net new memories
    • + *
+ * + * @param input the strongly-typed distillation input, injected by the Embabel framework + * @param context the action context providing access to AI capabilities + * @return a {@link MemoryDistillationResult} with memories to store and revoke + */ + @Action(description = "Distill relevant memories from a completed agent process") + @AchievesGoal(description = "Distill and manage memories from completed agent processes") + public MemoryDistillationResult distillMemories(MemoryDistillationInput input, ActionContext context) { + int currentCount = input.existingMemories().size(); + int capacityLeft = Math.max(0, input.maxTotalMemories() - currentCount); + int effectiveLimit = Math.min(capacityLeft, input.maxMemoriesAdded()); + + String prompt = buildPrompt(input.agentTask(), input.history(), input.blackboardObjects(), + input.existingMemories(), effectiveLimit); + + return context.ai() + .withDefaultLlm() + .createObject(prompt, MemoryDistillationResult.class); + } + + private String buildPrompt(Object agentTask, + List history, + List blackboardObjects, + List existingMemories, + int maxNewMemories) { + var sb = new StringBuilder(); + + sb.append(""" + You are a memory distillation agent. Analyze the following completed agent process \ + and determine which facts should be stored as new memories and which existing \ + memories should be revoked (because they are outdated, incorrect, or superseded). + + CONSTRAINTS: + """); + + sb.append("- The net number of new memories (stored count minus revoked count) must be <= ") + .append(maxNewMemories).append("\n"); + + sb.append(""" + - Only store memories that are actionable, likely to remain relevant, and \ + cannot always be inferred from limited context + - Only revoke memories that are clearly outdated, incorrect, or superseded \ + by new information + + """); + + sb.append("AGENT TASK:\n```json\n"); + if (agentTask != null) { + sb.append(serialize(agentTask)).append("\n"); + } else { + sb.append("null\n"); + } + sb.append("```\n"); + + sb.append("\nPROCESS HISTORY (actions executed):\n```json\n"); + if (history != null && !history.isEmpty()) { + sb.append("[\n"); + for (int i = 0; i < history.size(); i++) { + sb.append(serialize(history.get(i))); + if (i < history.size() - 1) sb.append(","); + sb.append("\n"); + } + sb.append("]\n"); + } else { + sb.append("[]\n"); + } + sb.append("```\n"); + + sb.append("\nBLACKBOARD OBJECTS (process context and results):\n```json\n"); + if (blackboardObjects != null && !blackboardObjects.isEmpty()) { + sb.append("[\n"); + for (int i = 0; i < blackboardObjects.size(); i++) { + sb.append(serialize(blackboardObjects.get(i))); + if (i < blackboardObjects.size() - 1) sb.append(","); + sb.append("\n"); + } + sb.append("]\n"); + } else { + sb.append("[]\n"); + } + sb.append("```\n"); + + sb.append("\nEXISTING MEMORIES:\n```json\n"); + if (existingMemories != null && !existingMemories.isEmpty()) { + sb.append("[\n"); + for (int i = 0; i < existingMemories.size(); i++) { + sb.append(serialize(existingMemories.get(i))); + if (i < existingMemories.size() - 1) sb.append(","); + sb.append("\n"); + } + sb.append("]\n"); + } else { + sb.append("[]\n"); + } + sb.append("```\n"); + + return sb.toString(); + } + + /** + * Serializes an object to JSON using the injected {@link ObjectMapper}. + * Falls back to {@link Object#toString()} if serialization fails. + */ + private String serialize(Object obj) { + try { + return objectMapper.writeValueAsString(obj); + } catch (Exception e) { + return obj.toString(); + } + } +} diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/events/MemoryRevokedEvent.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/events/MemoryRevokedEvent.java index c3599c94..8330c2e5 100644 --- a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/events/MemoryRevokedEvent.java +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/events/MemoryRevokedEvent.java @@ -31,7 +31,7 @@ *

This event is produced internally by the Embabel layer (via the agent's memory * management tools) when a memory is explicitly revoked, or as part of the * sliding-window eviction mechanism when - * {@link org.elasticsoftware.akces.annotations.AgenticAggregateInfo#maxMemories()} is exceeded. + * {@link org.elasticsoftware.akces.annotations.AgenticAggregateInfo#maxTotalMemories()} is exceeded. * * @param agenticAggregateId the unique identifier of the AgenticAggregate instance * @param memoryId UUID of the memory entry that was revoked diff --git a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntime.java b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntime.java index 728ad474..2edc1adc 100644 --- a/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntime.java +++ b/main/agentic/src/main/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntime.java @@ -17,12 +17,17 @@ package org.elasticsoftware.akces.agentic.runtime; +import com.embabel.agent.core.Agent; import com.embabel.agent.core.AgentPlatform; import com.embabel.agent.core.AgentProcess; import com.embabel.agent.core.AgentProcessStatusCode; +import com.embabel.agent.core.ProcessOptions; import jakarta.annotation.Nullable; import org.apache.kafka.common.errors.SerializationException; import org.elasticsoftware.akces.agentic.AgenticAggregateRuntime; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillationInput; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillationResult; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillerAgent; import org.elasticsoftware.akces.agentic.events.AgentTaskAssignedEvent; import org.elasticsoftware.akces.agentic.events.AgentTaskFinishedEvent; import org.elasticsoftware.akces.agentic.events.MemoryRevokedEvent; @@ -44,9 +49,7 @@ import java.io.IOException; import java.time.Instant; -import java.util.Collection; -import java.util.List; -import java.util.Objects; +import java.util.*; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.function.Consumer; @@ -71,6 +74,8 @@ public class KafkaAgenticAggregateRuntime implements AgenticAggregateRuntime { private final Class stateClass; private final AgentPlatform agentPlatform; private final AgenticAggregate aggregate; + private final int maxTotalMemories; + private final int maxMemoriesAdded; /** Round-robin counter for selecting the next agent task to resume. */ private final AtomicInteger nextTaskIndex = new AtomicInteger(0); @@ -78,25 +83,37 @@ public class KafkaAgenticAggregateRuntime implements AgenticAggregateRuntime { /** * Creates a new {@code KafkaAgenticAggregateRuntime}. * - * @param delegate the underlying aggregate runtime to delegate to - * @param objectMapper the Jackson {@link ObjectMapper} used for JSON serialization - * @param stateClass the aggregate state class - * @param agentPlatform the Embabel {@link AgentPlatform} used for AI-assisted processing; - * must not be {@code null} - * @param aggregate the agentic aggregate instance whose - * {@link AgenticAggregate#getCreateDomainEvent()} method provides the - * auto-create event + * @param delegate the underlying aggregate runtime to delegate to + * @param objectMapper the Jackson {@link ObjectMapper} used for JSON serialization + * @param stateClass the aggregate state class + * @param agentPlatform the Embabel {@link AgentPlatform} used for AI-assisted processing; + * must not be {@code null} + * @param aggregate the agentic aggregate instance whose + * {@link AgenticAggregate#getCreateDomainEvent()} method provides the + * auto-create event + * @param maxTotalMemories the total memory capacity for this aggregate + * @param maxMemoriesAdded the per-distillation budget for net new memories */ public KafkaAgenticAggregateRuntime(KafkaAggregateRuntime delegate, ObjectMapper objectMapper, Class stateClass, AgentPlatform agentPlatform, - AgenticAggregate aggregate) { + AgenticAggregate aggregate, + int maxTotalMemories, + int maxMemoriesAdded) { this.delegate = Objects.requireNonNull(delegate, "delegate must not be null"); this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null"); this.stateClass = Objects.requireNonNull(stateClass, "stateClass must not be null"); this.agentPlatform = Objects.requireNonNull(agentPlatform, "agentPlatform must not be null"); this.aggregate = Objects.requireNonNull(aggregate, "aggregate must not be null"); + if (maxTotalMemories < 0) { + throw new IllegalArgumentException("maxTotalMemories must be >= 0"); + } + if (maxMemoriesAdded < 0) { + throw new IllegalArgumentException("maxMemoriesAdded must be >= 0"); + } + this.maxTotalMemories = maxTotalMemories; + this.maxMemoriesAdded = maxMemoriesAdded; } // ------------------------------------------------------------------------- @@ -511,8 +528,171 @@ public void resumeNextAgentTask(Consumer protocolRecordConsumer, statusCode, Instant.now()); tickEvents = Stream.concat(tickEvents, Stream.of(finishedEvent)); + + // Distill memories from successfully completed processes + if (statusCode == AgentProcessStatusCode.COMPLETED) { + List memoryEvents = distillMemories(agentProcess, state, task); + if (!memoryEvents.isEmpty()) { + tickEvents = Stream.concat(tickEvents, memoryEvents.stream()); + } + } } delegate.processDomainEvents(tickEvents, task.agentProcessId(), protocolRecordConsumer, stateRecordSupplier); } + + // ------------------------------------------------------------------------- + // Memory distillation + // ------------------------------------------------------------------------- + + /** + * Distills relevant memories from a successfully completed {@link AgentProcess} + * by running the {@link MemoryDistillerAgent} to completion. + * + *

Creates a separate agent process for the {@link MemoryDistillerAgent}, populates + * its blackboard with a single {@link MemoryDistillationInput} containing the completed + * process's history, blackboard objects, current memories, and memory capacity constraints. + * The Embabel framework injects this input into the agent's action method via type-based + * parameter resolution. The process is then run to completion via {@link AgentProcess#run()}. + * + *

The net memory constraint ensures that + * {@code storedCount - revokedCount <= maxMemories - currentMemoryCount}, preventing + * the memory system from exceeding its configured capacity. + * + * @param completedProcess the agent process that has completed successfully + * @param state the current aggregate state + * @param task the assigned task that triggered the agent process + * @return a list of {@link MemoryStoredEvent} and {@link MemoryRevokedEvent} instances; + * may be empty if no memories need to be changed + */ + private List distillMemories(AgentProcess completedProcess, AggregateState state, + AssignedTask task) { + Agent memoryDistillerAgent = resolveMemoryDistillerAgent(); + if (memoryDistillerAgent == null) { + logger.warn("MemoryDistillerAgent not deployed on the platform; skipping memory distillation"); + return List.of(); + } + + List currentMemories = state instanceof MemoryAwareState mas + ? mas.getMemories() + : List.of(); + + MemoryDistillationInput distillationInput = new MemoryDistillationInput( + task, + completedProcess.getHistory(), + completedProcess.getBlackboard().getObjects(), + currentMemories, + maxTotalMemories, + maxMemoriesAdded); + + Map bindings = new LinkedHashMap<>(); + bindings.put("input", distillationInput); + + try { + AgentProcess distillerProcess = agentPlatform.createAgentProcess( + memoryDistillerAgent, ProcessOptions.DEFAULT, bindings); + distillerProcess.run(); + + MemoryDistillationResult result = distillerProcess.getBlackboard() + .last(MemoryDistillationResult.class); + + if (result == null) { + logger.debug("MemoryDistillerAgent produced no result for aggregate {}", getName()); + return List.of(); + } + + int capacityLeft = Math.max(0, maxTotalMemories - currentMemories.size()); + int effectiveLimit = Math.min(capacityLeft, maxMemoriesAdded); + return translateDistillationResult(result, currentMemories, effectiveLimit); + } catch (Exception e) { + logger.warn("Memory distillation failed for aggregate {}; proceeding without memory updates", + getName(), e); + return List.of(); + } + } + + /** + * Resolves the {@link MemoryDistillerAgent} from the {@link AgentPlatform}. + * + * @return the resolved agent, or {@code null} if not deployed + */ + private Agent resolveMemoryDistillerAgent() { + for (Agent agent : agentPlatform.agents()) { + if (MemoryDistillerAgent.AGENT_NAME.equals(agent.getName())) { + return agent; + } + } + return null; + } + + /** + * Translates a {@link MemoryDistillationResult} into a list of domain events, + * enforcing the net memory limit. + * + *

The constraint {@code stored.size() - revoked.size() <= maxNewMemories} is + * enforced by truncating stored memories when the limit would be exceeded. + * + * @param result the distillation result from the agent + * @param currentMemories the current memories from the aggregate state + * @param maxNewMemories the maximum number of net new memories allowed + * @return an unmodifiable list of domain events + */ + private List translateDistillationResult(MemoryDistillationResult result, + List currentMemories, + int maxNewMemories) { + List events = new ArrayList<>(); + + // Collect valid revocations first (only for memories that actually exist), + // de-duplicated by memoryId so capacity calculations match effective revocations. + Set currentMemoryIds = new HashSet<>(); + for (AgenticAggregateMemory mem : currentMemories) { + currentMemoryIds.add(mem.memoryId()); + } + + List validRevocations = new ArrayList<>(); + Set revokedMemoryIds = new HashSet<>(); + if (result.revoked() != null) { + for (MemoryRevokedEvent revoked : result.revoked()) { + String memoryId = revoked.memoryId(); + if (memoryId == null || !currentMemoryIds.contains(memoryId)) { + logger.debug("Skipping revocation of non-existent memory '{}' for aggregate {}", + memoryId, getName()); + continue; + } + if (!revokedMemoryIds.add(memoryId)) { + logger.debug("Skipping duplicate revocation of memory '{}' for aggregate {}", + memoryId, getName()); + continue; + } + validRevocations.add(revoked); + } + } + + // Calculate how many stored memories we can accept + int revokedCount = revokedMemoryIds.size(); + int maxStoredCount = maxNewMemories + revokedCount; + + // Add revocation events + events.addAll(validRevocations); + + // Add stored events (respecting the limit) + if (result.stored() != null) { + int storedCount = 0; + for (MemoryStoredEvent stored : result.stored()) { + if (storedCount >= maxStoredCount) { + logger.debug("Memory distillation limit reached ({} stored, {} revoked, max {}); " + + "truncating remaining stored memories for aggregate {}", + storedCount, revokedCount, maxStoredCount, getName()); + break; + } + events.add(stored); + storedCount++; + } + } + + logger.debug("Memory distillation for aggregate {} produced {} stored and {} revoked events", + getName(), events.size() - validRevocations.size(), validRevocations.size()); + + return List.copyOf(events); + } } diff --git a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResultTest.java b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResultTest.java new file mode 100644 index 00000000..4da92c90 --- /dev/null +++ b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/embabel/MemoryDistillationResultTest.java @@ -0,0 +1,61 @@ +/* + * Copyright 2022 - 2026 The Original Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.elasticsoftware.akces.agentic.embabel; + +import org.elasticsoftware.akces.agentic.events.MemoryRevokedEvent; +import org.elasticsoftware.akces.agentic.events.MemoryStoredEvent; +import org.junit.jupiter.api.Test; + +import java.time.Instant; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link MemoryDistillationResult}. + */ +class MemoryDistillationResultTest { + + private static final Instant NOW = Instant.now(); + + @Test + void shouldCreateWithStoredAndRevokedEvents() { + var stored = List.of( + new MemoryStoredEvent("agg-1", "mem-1", "testing", "Use JUnit 5", + "src/test", "best practice", NOW), + new MemoryStoredEvent("agg-1", "mem-2", "logging", "Use SLF4J", + "src/main", "consistency", NOW) + ); + var revoked = List.of( + new MemoryRevokedEvent("agg-1", "mem-old-1", "superseded by new info", NOW) + ); + + var result = new MemoryDistillationResult(stored, revoked); + + assertThat(result.stored()).hasSize(2); + assertThat(result.revoked()).hasSize(1); + } + + @Test + void shouldCreateWithEmptyLists() { + var result = new MemoryDistillationResult(List.of(), List.of()); + + assertThat(result.stored()).isEmpty(); + assertThat(result.revoked()).isEmpty(); + } +} diff --git a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/AgenticAggregateAutoCreateTest.java b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/AgenticAggregateAutoCreateTest.java index 2baa6ab4..db4172c2 100644 --- a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/AgenticAggregateAutoCreateTest.java +++ b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/AgenticAggregateAutoCreateTest.java @@ -240,7 +240,7 @@ void initializeStateShouldCallGetCreateDomainEventAndDelegateToRuntime() throws ObjectMapper objectMapper = JsonMapper.builder().build(); KafkaAgenticAggregateRuntime runtime = new KafkaAgenticAggregateRuntime( - delegate, objectMapper, TestBotState.class, agentPlatform, testBot); + delegate, objectMapper, TestBotState.class, agentPlatform, testBot, 100, 10); when(delegate.getName()).thenReturn("TestBot"); @@ -275,7 +275,7 @@ public Class getStateClass() { ObjectMapper objectMapper = JsonMapper.builder().build(); KafkaAgenticAggregateRuntime runtime = new KafkaAgenticAggregateRuntime( - delegate, objectMapper, TestBotState.class, agentPlatform, nullReturningAggregate); + delegate, objectMapper, TestBotState.class, agentPlatform, nullReturningAggregate, 100, 10); assertThatNullPointerException() .isThrownBy(() -> runtime.initializeState(pr -> {}, (der, ip) -> {})) @@ -289,7 +289,7 @@ void initializeStateShouldPassConsumersThroughToDelegate() throws IOException { ObjectMapper objectMapper = JsonMapper.builder().build(); KafkaAgenticAggregateRuntime runtime = new KafkaAgenticAggregateRuntime( - delegate, objectMapper, TestBotState.class, agentPlatform, testBot); + delegate, objectMapper, TestBotState.class, agentPlatform, testBot, 100, 10); when(delegate.getName()).thenReturn("TestBot"); diff --git a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntimeTest.java b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntimeTest.java index eb6cf968..af435135 100644 --- a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntimeTest.java +++ b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/KafkaAgenticAggregateRuntimeTest.java @@ -108,7 +108,7 @@ public String getAggregateId() { @BeforeEach void setUp() { objectMapper = JsonMapper.builder().build(); - runtime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, TestMemoryState.class, agentPlatform, aggregate); + runtime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, TestMemoryState.class, agentPlatform, aggregate, 100, 10); } // ------------------------------------------------------------------------- @@ -145,7 +145,7 @@ void getMemoriesShouldReturnMemoriesFromMemoryAwareState() throws IOException { @Test void getMemoriesShouldReturnEmptyListForNonMemoryAwareState() throws IOException { - var plainRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, PlainState.class, agentPlatform, aggregate); + var plainRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, PlainState.class, agentPlatform, aggregate, 100, 10); var state = new PlainState("agg-1"); byte[] payload = objectMapper.writeValueAsBytes(state); @@ -235,7 +235,7 @@ void getAgentPlatformShouldReturnInjectedPlatform() { void constructorShouldRejectNullDelegate() { assertThatNullPointerException() .isThrownBy(() -> new KafkaAgenticAggregateRuntime( - null, objectMapper, TestMemoryState.class, agentPlatform, aggregate)) + null, objectMapper, TestMemoryState.class, agentPlatform, aggregate, 100, 10)) .withMessageContaining("delegate"); } @@ -243,7 +243,7 @@ void constructorShouldRejectNullDelegate() { void constructorShouldRejectNullObjectMapper() { assertThatNullPointerException() .isThrownBy(() -> new KafkaAgenticAggregateRuntime( - delegate, null, TestMemoryState.class, agentPlatform, aggregate)) + delegate, null, TestMemoryState.class, agentPlatform, aggregate, 100, 10)) .withMessageContaining("objectMapper"); } @@ -251,7 +251,7 @@ void constructorShouldRejectNullObjectMapper() { void constructorShouldRejectNullStateClass() { assertThatNullPointerException() .isThrownBy(() -> new KafkaAgenticAggregateRuntime( - delegate, objectMapper, null, agentPlatform, aggregate)) + delegate, objectMapper, null, agentPlatform, aggregate, 100, 10)) .withMessageContaining("stateClass"); } @@ -259,7 +259,7 @@ void constructorShouldRejectNullStateClass() { void constructorShouldRejectNullAgentPlatform() { assertThatNullPointerException() .isThrownBy(() -> new KafkaAgenticAggregateRuntime( - delegate, objectMapper, TestMemoryState.class, null, aggregate)) + delegate, objectMapper, TestMemoryState.class, null, aggregate, 100, 10)) .withMessageContaining("agentPlatform"); } @@ -267,7 +267,7 @@ void constructorShouldRejectNullAgentPlatform() { void constructorShouldRejectNullAggregate() { assertThatNullPointerException() .isThrownBy(() -> new KafkaAgenticAggregateRuntime( - delegate, objectMapper, TestMemoryState.class, agentPlatform, null)) + delegate, objectMapper, TestMemoryState.class, agentPlatform, null, 100, 10)) .withMessageContaining("aggregate"); } diff --git a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/MemoryDistillationTest.java b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/MemoryDistillationTest.java new file mode 100644 index 00000000..0a6fe129 --- /dev/null +++ b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/MemoryDistillationTest.java @@ -0,0 +1,654 @@ +/* + * Copyright 2022 - 2026 The Original Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.elasticsoftware.akces.agentic.runtime; + +import com.embabel.agent.core.Agent; +import com.embabel.agent.core.AgentPlatform; +import com.embabel.agent.core.AgentProcess; +import com.embabel.agent.core.AgentProcessStatusCode; +import com.embabel.agent.core.Blackboard; +import com.embabel.agent.core.ProcessOptions; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillationInput; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillationResult; +import org.elasticsoftware.akces.agentic.embabel.MemoryDistillerAgent; +import org.elasticsoftware.akces.agentic.events.AgentTaskFinishedEvent; +import org.elasticsoftware.akces.agentic.events.MemoryRevokedEvent; +import org.elasticsoftware.akces.agentic.events.MemoryStoredEvent; +import org.elasticsoftware.akces.aggregate.*; +import org.elasticsoftware.akces.commands.CommandBus; +import org.elasticsoftware.akces.events.DomainEvent; +import org.elasticsoftware.akces.kafka.KafkaAggregateRuntime; +import org.elasticsoftware.akces.protocol.AggregateStateRecord; +import org.elasticsoftware.akces.protocol.PayloadEncoding; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.json.JsonMapper; + +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +/** + * Unit tests for the memory distillation integration in + * {@link KafkaAgenticAggregateRuntime#resumeNextAgentTask}. + */ +@ExtendWith(MockitoExtension.class) +class MemoryDistillationTest { + + /** State that implements both TaskAwareState and MemoryAwareState. */ + record MemoryTaskState( + String id, + List assignedTasks, + List memories + ) implements AggregateState, TaskAwareState, MemoryAwareState { + + @Override + public String getAggregateId() { + return id; + } + + @Override + public List getAssignedTasks() { + return assignedTasks; + } + + @Override + public TaskAwareState withAssignedTask(AssignedTask task) { + var newTasks = new ArrayList<>(assignedTasks); + newTasks.add(task); + return new MemoryTaskState(id, List.copyOf(newTasks), memories); + } + + @Override + public TaskAwareState withoutAssignedTask(String agentProcessId) { + return new MemoryTaskState(id, assignedTasks.stream() + .filter(t -> !t.agentProcessId().equals(agentProcessId)) + .toList(), memories); + } + + @Override + public List getMemories() { + return memories; + } + + @Override + public MemoryAwareState withMemory(AgenticAggregateMemory memory) { + var updated = new ArrayList<>(memories); + updated.add(memory); + return new MemoryTaskState(id, assignedTasks, List.copyOf(updated)); + } + + @Override + public MemoryAwareState withoutMemory(String memoryId) { + return new MemoryTaskState(id, assignedTasks, memories.stream() + .filter(m -> !m.memoryId().equals(memoryId)) + .toList()); + } + } + + private static final Instant NOW = Instant.now(); + + @Mock + private KafkaAggregateRuntime delegate; + + @Mock + private AgentPlatform agentPlatform; + + @Mock + private AgentProcess agentProcess; + + @Mock + private AgentProcess distillerProcess; + + @Mock + private Blackboard blackboard; + + @Mock + private Blackboard distillerBlackboard; + + @Mock + private CommandBus commandBus; + + @Mock + private AgenticAggregate aggregate; + + @Mock + private Agent memoryDistillerAgent; + + private ObjectMapper objectMapper; + private KafkaAgenticAggregateRuntime runtime; + + @BeforeEach + void setUp() { + objectMapper = JsonMapper.builder().build(); + runtime = new KafkaAgenticAggregateRuntime( + delegate, objectMapper, MemoryTaskState.class, agentPlatform, aggregate, 100, 10); + } + + // ------------------------------------------------------------------------- + // Memory distillation on COMPLETED process + // ------------------------------------------------------------------------- + + @Test + @SuppressWarnings("unchecked") + void resumeShouldDistillMemoriesOnCompletedProcess() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + // Process finishes with COMPLETED status + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + // Set up MemoryDistillerAgent resolution + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Set up distiller process + var distillationResult = new MemoryDistillationResult( + List.of(new MemoryStoredEvent("agg-1", "mem-new-1", "testing", "Use JUnit 5", + "src/test", "best practice", NOW)), + List.of()); + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(distillationResult); + + // Capture the events stream + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // Should have: tick events (empty) + AgentTaskFinishedEvent + MemoryStoredEvent + assertThat(events).hasSize(2); + assertThat(events.get(0)).isInstanceOf(AgentTaskFinishedEvent.class); + assertThat(events.get(1)).isInstanceOf(MemoryStoredEvent.class); + + MemoryStoredEvent storedEvent = (MemoryStoredEvent) events.get(1); + assertThat(storedEvent.agenticAggregateId()).isEqualTo("agg-1"); + assertThat(storedEvent.subject()).isEqualTo("testing"); + assertThat(storedEvent.fact()).isEqualTo("Use JUnit 5"); + + // Verify distiller process was run to completion + verify(distillerProcess).run(); + } + + @Test + @SuppressWarnings("unchecked") + void resumeShouldNotDistillMemoriesOnFailedProcess() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + // Process finishes with FAILED status + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.FAILED); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // Should only have AgentTaskFinishedEvent, no memory distillation + assertThat(events).hasSize(1); + assertThat(events.getFirst()).isInstanceOf(AgentTaskFinishedEvent.class); + + // Should NOT create a distiller process + verify(agentPlatform, never()).createAgentProcess( + argThat(a -> MemoryDistillerAgent.AGENT_NAME.equals(a.getName())), + any(), anyMap()); + } + + @Test + @SuppressWarnings("unchecked") + void resumeShouldSkipDistillationWhenDistillerAgentNotDeployed() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + + // No MemoryDistillerAgent deployed + when(agentPlatform.agents()).thenReturn(List.of()); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // Should only have AgentTaskFinishedEvent + assertThat(events).hasSize(1); + assertThat(events.getFirst()).isInstanceOf(AgentTaskFinishedEvent.class); + } + + @Test + @SuppressWarnings("unchecked") + void resumeShouldHandleDistillationFailureGracefully() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Distiller throws an exception + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenThrow(new RuntimeException("LLM unavailable")); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // Should still emit AgentTaskFinishedEvent despite distillation failure + assertThat(events).hasSize(1); + assertThat(events.getFirst()).isInstanceOf(AgentTaskFinishedEvent.class); + } + + // ------------------------------------------------------------------------- + // Memory limit enforcement + // ------------------------------------------------------------------------- + + @Test + @SuppressWarnings("unchecked") + void distillationShouldEnforceMaxMemoriesLimit() throws IOException { + // maxTotalMemories = 5, maxMemoriesAdded = 10 — capacity is the binding constraint + runtime = new KafkaAgenticAggregateRuntime( + delegate, objectMapper, MemoryTaskState.class, agentPlatform, aggregate, 5, 10); + + // State already has 3 memories, so capacityLeft = 2, effectiveLimit = min(2, 10) = 2 + var existingMemories = List.of( + new AgenticAggregateMemory("mem-1", "s1", "f1", "c1", "r1", Instant.now()), + new AgenticAggregateMemory("mem-2", "s2", "f2", "c2", "r2", Instant.now()), + new AgenticAggregateMemory("mem-3", "s3", "f3", "c3", "r3", Instant.now()) + ); + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), existingMemories); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Agent tries to store 5 memories but only 2 should be allowed + var distillationResult = new MemoryDistillationResult( + List.of( + new MemoryStoredEvent("agg-1", "n1", "s1", "fact1", "c1", "r1", NOW), + new MemoryStoredEvent("agg-1", "n2", "s2", "fact2", "c2", "r2", NOW), + new MemoryStoredEvent("agg-1", "n3", "s3", "fact3", "c3", "r3", NOW), + new MemoryStoredEvent("agg-1", "n4", "s4", "fact4", "c4", "r4", NOW), + new MemoryStoredEvent("agg-1", "n5", "s5", "fact5", "c5", "r5", NOW) + ), + List.of()); + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(distillationResult); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // 1 finished + 2 stored (truncated from 5) + assertThat(events).hasSize(3); + assertThat(events.get(0)).isInstanceOf(AgentTaskFinishedEvent.class); + long storedCount = events.stream().filter(e -> e instanceof MemoryStoredEvent).count(); + assertThat(storedCount).isEqualTo(2); + } + + @Test + @SuppressWarnings("unchecked") + void distillationShouldEnforceMaxMemoriesAddedLimit() throws IOException { + // maxTotalMemories = 100, maxMemoriesAdded = 2 — per-distillation budget is the constraint + runtime = new KafkaAgenticAggregateRuntime( + delegate, objectMapper, MemoryTaskState.class, agentPlatform, aggregate, 100, 2); + + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Agent tries to store 5 memories but only 2 should be allowed (maxMemoriesAdded=2) + var distillationResult = new MemoryDistillationResult( + List.of( + new MemoryStoredEvent("agg-1", "n1", "s1", "fact1", "c1", "r1", NOW), + new MemoryStoredEvent("agg-1", "n2", "s2", "fact2", "c2", "r2", NOW), + new MemoryStoredEvent("agg-1", "n3", "s3", "fact3", "c3", "r3", NOW), + new MemoryStoredEvent("agg-1", "n4", "s4", "fact4", "c4", "r4", NOW), + new MemoryStoredEvent("agg-1", "n5", "s5", "fact5", "c5", "r5", NOW) + ), + List.of()); + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(distillationResult); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // 1 finished + 2 stored (truncated from 5 by maxMemoriesAdded) + assertThat(events).hasSize(3); + assertThat(events.get(0)).isInstanceOf(AgentTaskFinishedEvent.class); + long storedCount = events.stream().filter(e -> e instanceof MemoryStoredEvent).count(); + assertThat(storedCount).isEqualTo(2); + } + + @Test + @SuppressWarnings("unchecked") + void distillationShouldAllowMoreStoredWhenMemoriesAreRevoked() throws IOException { + // maxTotalMemories = 5, maxMemoriesAdded = 10; existing = 4, so capacityLeft = 1 + // But we also revoke 2, so we can store up to 3 (1 + 2) + runtime = new KafkaAgenticAggregateRuntime( + delegate, objectMapper, MemoryTaskState.class, agentPlatform, aggregate, 5, 10); + + var existingMemories = List.of( + new AgenticAggregateMemory("mem-1", "s1", "f1", "c1", "r1", Instant.now()), + new AgenticAggregateMemory("mem-2", "s2", "f2", "c2", "r2", Instant.now()), + new AgenticAggregateMemory("mem-3", "s3", "f3", "c3", "r3", Instant.now()), + new AgenticAggregateMemory("mem-4", "s4", "f4", "c4", "r4", Instant.now()) + ); + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), existingMemories); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Revoke 2, store 3 + var distillationResult = new MemoryDistillationResult( + List.of( + new MemoryStoredEvent("agg-1", "n1", "s1", "new-fact1", "c1", "r1", NOW), + new MemoryStoredEvent("agg-1", "n2", "s2", "new-fact2", "c2", "r2", NOW), + new MemoryStoredEvent("agg-1", "n3", "s3", "new-fact3", "c3", "r3", NOW) + ), + List.of( + new MemoryRevokedEvent("agg-1", "mem-1", "outdated", NOW), + new MemoryRevokedEvent("agg-1", "mem-2", "superseded", NOW) + )); + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(distillationResult); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // 1 finished + 2 revoked + 3 stored = 6 + assertThat(events).hasSize(6); + assertThat(events.get(0)).isInstanceOf(AgentTaskFinishedEvent.class); + + long revokedCount = events.stream().filter(e -> e instanceof MemoryRevokedEvent).count(); + long storedCount = events.stream().filter(e -> e instanceof MemoryStoredEvent).count(); + assertThat(revokedCount).isEqualTo(2); + assertThat(storedCount).isEqualTo(3); + } + + @Test + @SuppressWarnings("unchecked") + void distillationShouldSkipRevocationOfNonExistentMemories() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var existingMemories = List.of( + new AgenticAggregateMemory("mem-1", "s1", "f1", "c1", "r1", Instant.now()) + ); + var state = new MemoryTaskState("agg-1", List.of(task), existingMemories); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Try to revoke a memory that doesn't exist + var distillationResult = new MemoryDistillationResult( + List.of(), + List.of( + new MemoryRevokedEvent("agg-1", "mem-1", "outdated", NOW), + new MemoryRevokedEvent("agg-1", "non-existent", "cleanup", NOW) + )); + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(distillationResult); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // 1 finished + 1 revoked (only mem-1, not non-existent) + long revokedCount = events.stream().filter(e -> e instanceof MemoryRevokedEvent).count(); + assertThat(revokedCount).isEqualTo(1); + + MemoryRevokedEvent revokedEvent = events.stream() + .filter(e -> e instanceof MemoryRevokedEvent) + .map(e -> (MemoryRevokedEvent) e) + .findFirst() + .orElseThrow(); + assertThat(revokedEvent.memoryId()).isEqualTo("mem-1"); + } + + @Test + @SuppressWarnings("unchecked") + void distillationShouldHandleNullResultGracefully() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of()); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of()); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + // Distiller produces no result + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(null); + + ArgumentCaptor> eventsCaptor = ArgumentCaptor.forClass(Stream.class); + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(delegate).processDomainEvents(eventsCaptor.capture(), eq("proc-1"), any(), any()); + List events = eventsCaptor.getValue().toList(); + + // Should only have AgentTaskFinishedEvent + assertThat(events).hasSize(1); + assertThat(events.getFirst()).isInstanceOf(AgentTaskFinishedEvent.class); + } + + @Test + @SuppressWarnings("unchecked") + void distillationShouldPassCorrectBindingsToDistillerProcess() throws IOException { + var party = new HumanRequestingParty("user-1", "analyst"); + var task = new AssignedTask("proc-1", "Task", party, Map.of(), Instant.now()); + var existingMemory = new AgenticAggregateMemory( + "mem-1", "testing", "Use JUnit", "cite", "reason", Instant.now()); + var state = new MemoryTaskState("agg-1", List.of(task), List.of(existingMemory)); + byte[] payload = objectMapper.writeValueAsBytes(state); + var stateRecord = new AggregateStateRecord(null, "MemoryTaskState", 1, payload, + PayloadEncoding.JSON, "agg-1", null, 1L); + when(delegate.materializeState(stateRecord)).thenReturn(state); + + when(agentPlatform.getAgentProcess("proc-1")).thenReturn(agentProcess); + when(agentProcess.getBlackboard()).thenReturn(blackboard); + when(blackboard.getObjects()).thenReturn(List.of("some-object")); + when(delegate.getAllDomainEventTypes()).thenReturn(List.of()); + + when(agentProcess.getFinished()).thenReturn(true); + when(agentProcess.getStatus()).thenReturn(AgentProcessStatusCode.COMPLETED); + when(agentProcess.getHistory()).thenReturn(List.of()); + + when(memoryDistillerAgent.getName()).thenReturn(MemoryDistillerAgent.AGENT_NAME); + when(agentPlatform.agents()).thenReturn(List.of(memoryDistillerAgent)); + + when(agentPlatform.createAgentProcess(eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), anyMap())) + .thenReturn(distillerProcess); + when(distillerProcess.getBlackboard()).thenReturn(distillerBlackboard); + when(distillerBlackboard.last(MemoryDistillationResult.class)).thenReturn(null); + + @SuppressWarnings("rawtypes") + ArgumentCaptor> bindingsCaptor = ArgumentCaptor.forClass(Map.class); + + runtime.resumeNextAgentTask(pr -> {}, () -> stateRecord, commandBus); + + verify(agentPlatform).createAgentProcess( + eq(memoryDistillerAgent), eq(ProcessOptions.DEFAULT), bindingsCaptor.capture()); + + Map bindings = bindingsCaptor.getValue(); + assertThat(bindings).containsKey("input"); + assertThat(bindings.get("input")).isInstanceOf(MemoryDistillationInput.class); + + MemoryDistillationInput input = (MemoryDistillationInput) bindings.get("input"); + assertThat(input.agentTask()).isEqualTo(task); + assertThat(input.history()).isEqualTo(List.of()); + assertThat(input.blackboardObjects()).isEqualTo(List.of("some-object")); + assertThat(input.existingMemories()).isEqualTo(List.of(existingMemory)); + assertThat(input.maxTotalMemories()).isEqualTo(100); + assertThat(input.maxMemoriesAdded()).isEqualTo(10); + } +} diff --git a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/ResumeNextAgentTaskTest.java b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/ResumeNextAgentTaskTest.java index cb4d1b78..3f9855d4 100644 --- a/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/ResumeNextAgentTaskTest.java +++ b/main/agentic/src/test/java/org/elasticsoftware/akces/agentic/runtime/ResumeNextAgentTaskTest.java @@ -113,7 +113,7 @@ public String getAggregateId() { @BeforeEach void setUp() { objectMapper = JsonMapper.builder().build(); - runtime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, TaskState.class, agentPlatform, aggregate); + runtime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, TaskState.class, agentPlatform, aggregate, 100, 10); } @Test @@ -126,7 +126,7 @@ void resumeShouldDoNothingWhenStateIsNull() throws IOException { @Test void resumeShouldDoNothingWhenStateDoesNotImplementTaskAwareState() throws IOException { - var simpleRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, SimpleState.class, agentPlatform, aggregate); + var simpleRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, SimpleState.class, agentPlatform, aggregate, 100, 10); var state = new SimpleState("agg-1"); byte[] payload = objectMapper.writeValueAsBytes(state); var stateRecord = new AggregateStateRecord(null, "SimpleState", 1, payload, @@ -235,7 +235,7 @@ void hasActiveAgentTasksShouldReturnFalseWhenStateIsNull() throws IOException { @Test void hasActiveAgentTasksShouldReturnFalseWhenNotTaskAware() throws IOException { - var simpleRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, SimpleState.class, agentPlatform, aggregate); + var simpleRuntime = new KafkaAgenticAggregateRuntime(delegate, objectMapper, SimpleState.class, agentPlatform, aggregate, 100, 10); var state = new SimpleState("agg-1"); byte[] payload = objectMapper.writeValueAsBytes(state); var stateRecord = new AggregateStateRecord(null, "SimpleState", 1, payload, diff --git a/main/api/src/main/java/org/elasticsoftware/akces/annotations/AgenticAggregateInfo.java b/main/api/src/main/java/org/elasticsoftware/akces/annotations/AgenticAggregateInfo.java index dee3b6bd..252ffe40 100644 --- a/main/api/src/main/java/org/elasticsoftware/akces/annotations/AgenticAggregateInfo.java +++ b/main/api/src/main/java/org/elasticsoftware/akces/annotations/AgenticAggregateInfo.java @@ -65,7 +65,17 @@ * The sliding-window capacity of the memory system. When the number of stored memories * exceeds this limit the oldest entries are evicted to make room for new ones. */ - int maxMemories() default 100; + int maxTotalMemories() default 50; + + /** + * The maximum number of net new memories that the {@code MemoryDistillerAgent} may add + * in a single distillation pass (i.e. {@code stored − revoked ≤ maxMemoriesAdded}). + * + *

This acts as a per-distillation budget, independent of {@link #maxTotalMemories()}. + * Set to a low value to avoid overwhelming the memory store with many entries in a single + * agent run. + */ + int maxMemoriesAdded() default 5; /** * Command classes to be processed by the AI agent instead of a deterministic diff --git a/main/api/src/test/java/org/elasticsoftware/akces/aggregate/AgenticAggregateInfoTest.java b/main/api/src/test/java/org/elasticsoftware/akces/aggregate/AgenticAggregateInfoTest.java index f8538add..bd63f18b 100644 --- a/main/api/src/test/java/org/elasticsoftware/akces/aggregate/AgenticAggregateInfoTest.java +++ b/main/api/src/test/java/org/elasticsoftware/akces/aggregate/AgenticAggregateInfoTest.java @@ -66,8 +66,8 @@ public DomainEvent getCreateDomainEvent() { } } - /** Test aggregate annotated with a custom maxMemories value. */ - @AgenticAggregateInfo(value = "CustomAgentic", stateClass = TestState.class, maxMemories = 50) + /** Test aggregate annotated with custom maxTotalMemories and maxMemoriesAdded values. */ + @AgenticAggregateInfo(value = "CustomAgentic", stateClass = TestState.class, maxTotalMemories = 50, maxMemoriesAdded = 5) static class CustomMaxMemoriesAggregate implements AgenticAggregate { @Override public Class getStateClass() { @@ -81,17 +81,31 @@ public DomainEvent getCreateDomainEvent() { } @Test - void defaultMaxMemoriesShouldBe100() { + void defaultMaxMemoriesShouldBe50() { AgenticAggregateInfo info = DefaultAgenticAggregate.class.getAnnotation(AgenticAggregateInfo.class); assertThat(info).isNotNull(); - assertThat(info.maxMemories()).isEqualTo(100); + assertThat(info.maxTotalMemories()).isEqualTo(50); + } + + @Test + void defaultMaxMemoriesAddedShouldBe5() { + AgenticAggregateInfo info = DefaultAgenticAggregate.class.getAnnotation(AgenticAggregateInfo.class); + assertThat(info).isNotNull(); + assertThat(info.maxMemoriesAdded()).isEqualTo(5); } @Test void customMaxMemoriesShouldBeHonored() { AgenticAggregateInfo info = CustomMaxMemoriesAggregate.class.getAnnotation(AgenticAggregateInfo.class); assertThat(info).isNotNull(); - assertThat(info.maxMemories()).isEqualTo(50); + assertThat(info.maxTotalMemories()).isEqualTo(50); + } + + @Test + void customMaxMemoriesAddedShouldBeHonored() { + AgenticAggregateInfo info = CustomMaxMemoriesAggregate.class.getAnnotation(AgenticAggregateInfo.class); + assertThat(info).isNotNull(); + assertThat(info.maxMemoriesAdded()).isEqualTo(5); } @Test diff --git a/main/eventcatalog/src/test/java/org/elasticsoftware/akces/eventcatalog/EventCatalogProcessorTest.java b/main/eventcatalog/src/test/java/org/elasticsoftware/akces/eventcatalog/EventCatalogProcessorTest.java index 86222210..8892542c 100644 --- a/main/eventcatalog/src/test/java/org/elasticsoftware/akces/eventcatalog/EventCatalogProcessorTest.java +++ b/main/eventcatalog/src/test/java/org/elasticsoftware/akces/eventcatalog/EventCatalogProcessorTest.java @@ -1000,7 +1000,7 @@ public AssistantState withoutMemory(String memoryId) { value = "Assistant", stateClass = AssistantState.class, description = "AI Assistant AgenticAggregate", - maxMemories = 50) + maxTotalMemories = 50) @SuppressWarnings("unused") public final class Assistant implements AgenticAggregate { @Override