diff --git a/sentinel-ai-agent-memory/src/main/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutput.java b/sentinel-ai-agent-memory/src/main/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutput.java index 877d284..76101d3 100644 --- a/sentinel-ai-agent-memory/src/main/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutput.java +++ b/sentinel-ai-agent-memory/src/main/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutput.java @@ -12,7 +12,7 @@ * */ @Value -@Builder +@Builder(builderClassName = "AgentMemoryOutputBuilder") @Jacksonized @JsonClassDescription("Output for memory extraction from conversations by the model") public class AgentMemoryOutput { @@ -20,4 +20,8 @@ public class AgentMemoryOutput { List of memories extracted by the LLM from the conversation. These memories will be saved and retrieved by later by the agent to gain more intelligence over time. """) List generatedMemory; + + public static class AgentMemoryOutputBuilder { + public AgentMemoryOutputBuilder() {} + } } diff --git a/sentinel-ai-agent-memory/src/test/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutputTest.java b/sentinel-ai-agent-memory/src/test/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutputTest.java new file mode 100644 index 0000000..3a7a427 --- /dev/null +++ b/sentinel-ai-agent-memory/src/test/java/com/phonepe/sentinelai/agentmemory/AgentMemoryOutputTest.java @@ -0,0 +1,185 @@ +package com.phonepe.sentinelai.agentmemory; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.phonepe.sentinelai.core.utils.JsonUtils; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Tests for {@link AgentMemoryOutput} and its custom builder + */ +@Slf4j +class AgentMemoryOutputTest { + + private final ObjectMapper objectMapper = JsonUtils.createMapper(); + + @Test + void testBuilderCreation() { + AgentMemoryOutput.AgentMemoryOutputBuilder builder = AgentMemoryOutput.builder(); + assertNotNull(builder); + } + + @Test + void testBuilderWithEmptyGeneratedMemory() { + // Test building with empty list + List emptyList = List.of(); + AgentMemoryOutput output = AgentMemoryOutput.builder() + .generatedMemory(emptyList) + .build(); + + assertNotNull(output); + assertNotNull(output.getGeneratedMemory()); + assertTrue(output.getGeneratedMemory().isEmpty()); + } + + @Test + void testBuilderWithGeneratedMemory() { + // Create sample memory units + GeneratedMemoryUnit memory1 = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user123", + MemoryType.SEMANTIC, + "UserName", + "The user's name is John Doe", + List.of("user", "personal"), + 8 + ); + + GeneratedMemoryUnit memory2 = new GeneratedMemoryUnit( + MemoryScope.AGENT, + "agent-proc", + MemoryType.PROCEDURAL, + "WeatherQuery", + "When user asks about weather, first check their location", + List.of("weather", "procedure"), + 7 + ); + + List memories = Arrays.asList(memory1, memory2); + + // Test building with generated memories + AgentMemoryOutput output = AgentMemoryOutput.builder() + .generatedMemory(memories) + .build(); + + assertNotNull(output); + assertNotNull(output.getGeneratedMemory()); + assertEquals(2, output.getGeneratedMemory().size()); + assertEquals(memories, output.getGeneratedMemory()); + } + + @Test + void testBuilderToString() { + // Test builder toString method + AgentMemoryOutput.AgentMemoryOutputBuilder builder = AgentMemoryOutput.builder(); + String builderString = builder.toString(); + + assertNotNull(builderString); + assertTrue(builderString.contains("AgentMemoryOutputBuilder")); + assertTrue(builderString.contains("generatedMemory=null")); + } + + @Test + void testBuilderEqualsAndHashCode() { + // Test that two empty outputs are equal + AgentMemoryOutput output1 = AgentMemoryOutput.builder().build(); + AgentMemoryOutput output2 = AgentMemoryOutput.builder().build(); + + assertEquals(output1, output2); + assertEquals(output1.hashCode(), output2.hashCode()); + } + + @Test + void testJsonSerialization() throws Exception { + // Test JSON serialization/deserialization + GeneratedMemoryUnit memory = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user123", + MemoryType.SEMANTIC, + "UserPreference", + "User prefers dark mode", + List.of("ui", "preference"), + 6 + ); + + AgentMemoryOutput original = AgentMemoryOutput.builder() + .generatedMemory(List.of(memory)) + .build(); + + // Serialize to JSON + String json = objectMapper.writeValueAsString(original); + assertNotNull(json); + assertTrue(json.contains("UserPreference")); + assertTrue(json.contains("dark mode")); + + // Deserialize back + AgentMemoryOutput deserialized = objectMapper.readValue(json, AgentMemoryOutput.class); + assertNotNull(deserialized); + assertEquals(original, deserialized); + } + + @Test + void testCustomBuilderClassName() { + // Test that the custom builder class name is correctly set + // This is important for Jackson deserialization + AgentMemoryOutput.AgentMemoryOutputBuilder builder = AgentMemoryOutput.builder(); + assertEquals("AgentMemoryOutputBuilder", builder.getClass().getSimpleName()); + } + + @Test + void testBuilderWithComplexMemoryStructure() { + // Test with a more complex memory structure + List memories = Arrays.asList( + new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user456", + MemoryType.SEMANTIC, + "PersonalInfo", + "User is a software engineer from India working on AI systems", + Arrays.asList("personal", "profession", "location", "technology"), + 9 + ), + new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user456", + MemoryType.EPISODIC, + "PreviousConversation", + "User asked about best practices for microservices architecture yesterday", + Arrays.asList("conversation", "microservices", "architecture"), + 7 + ), + new GeneratedMemoryUnit( + MemoryScope.AGENT, + "sentiment-analyzer", + MemoryType.PROCEDURAL, + "SentimentAnalysis", + "When analyzing sentiment, consider cultural context and technical jargon", + Arrays.asList("sentiment", "analysis", "cultural", "technical"), + 8 + ) + ); + + AgentMemoryOutput output = AgentMemoryOutput.builder() + .generatedMemory(memories) + .build(); + + assertNotNull(output); + assertEquals(3, output.getGeneratedMemory().size()); + + // Verify different memory types and scopes are preserved + long entityMemories = output.getGeneratedMemory().stream() + .filter(m -> m.getScope() == MemoryScope.ENTITY) + .count(); + long agentMemories = output.getGeneratedMemory().stream() + .filter(m -> m.getScope() == MemoryScope.AGENT) + .count(); + + assertEquals(2, entityMemories); + assertEquals(1, agentMemories); + } +} diff --git a/sentinel-ai-storage-es/src/main/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorage.java b/sentinel-ai-storage-es/src/main/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorage.java new file mode 100644 index 0000000..d1bfe60 --- /dev/null +++ b/sentinel-ai-storage-es/src/main/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorage.java @@ -0,0 +1,38 @@ +package com.phonepe.sentinelai.storage.memory; + +import com.phonepe.sentinelai.agentmemory.AgentMemory; +import com.phonepe.sentinelai.agentmemory.AgentMemoryStore; +import com.phonepe.sentinelai.agentmemory.MemoryScope; +import com.phonepe.sentinelai.agentmemory.MemoryType; +import lombok.extern.slf4j.Slf4j; + +import java.util.*; +import java.util.concurrent.ConcurrentHashMap; + +@Slf4j +public class InMemoryMemoryStorage implements AgentMemoryStore { + + private record Key(MemoryScope scope, String scopeId) { + } + + private final Map> memories = new ConcurrentHashMap<>(); + + @Override + public List findMemories(String scopeId, MemoryScope scope, Set memoryTypes, List topics, String query, int minReusabilityScore, int count) { + if (scopeId == null || scope == null) { + return memories.values().stream() + .flatMap(List::stream) + .toList(); + } else { + return memories.getOrDefault(new Key(scope, scopeId), List.of()); + } + } + + @Override + public Optional save(AgentMemory agentMemory) { + final var key = new Key(agentMemory.getScope(), agentMemory.getScopeId()); + final var memoriesInScope = memories.computeIfAbsent(key, k -> new ArrayList<>()); + memoriesInScope.add(agentMemory); + return Optional.of(agentMemory); + } +} diff --git a/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/AgentMemoryOutputIntegrationTest.java b/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/AgentMemoryOutputIntegrationTest.java new file mode 100644 index 0000000..c348cc0 --- /dev/null +++ b/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/AgentMemoryOutputIntegrationTest.java @@ -0,0 +1,320 @@ +package com.phonepe.sentinelai.storage.memory; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.phonepe.sentinelai.agentmemory.*; +import com.phonepe.sentinelai.core.utils.JsonUtils; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Integration tests for AgentMemoryOutput builder working with InMemoryMemoryStorage + */ +@Slf4j +class AgentMemoryOutputIntegrationTest { + + private ObjectMapper objectMapper; + private InMemoryMemoryStorage memoryStorage; + + @BeforeEach + void setUp() { + objectMapper = JsonUtils.createMapper(); + memoryStorage = new InMemoryMemoryStorage(); + } + + @Test + void testAgentMemoryOutputBuilderWithInMemoryStorage() throws Exception { + // Create AgentMemoryOutput using the custom builder + GeneratedMemoryUnit memory1 = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user123", + MemoryType.SEMANTIC, + "UserName", + "The user's name is Alice Johnson", + List.of("personal", "identity"), + 8 + ); + + GeneratedMemoryUnit memory2 = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user123", + MemoryType.SEMANTIC, + "UserPreference", + "User prefers email notifications over SMS", + List.of("preferences", "communication"), + 7 + ); + + AgentMemoryOutput memoryOutput = AgentMemoryOutput.builder() + .generatedMemory(Arrays.asList(memory1, memory2)) + .build(); + + assertNotNull(memoryOutput); + assertEquals(2, memoryOutput.getGeneratedMemory().size()); + + // Convert memories to AgentMemory objects and save to storage + for (GeneratedMemoryUnit generatedMemory : memoryOutput.getGeneratedMemory()) { + AgentMemory agentMemory = convertToAgentMemory(generatedMemory, "test-agent"); + memoryStorage.save(agentMemory); + } + + // Verify memories were saved to InMemoryStorage + List savedMemories = memoryStorage.findMemories( + "user123", MemoryScope.ENTITY, null, null, null, 0, 10 + ); + + assertEquals(2, savedMemories.size()); + + // Check that the memories have correct content + boolean foundUserName = savedMemories.stream() + .anyMatch(m -> "UserName".equals(m.getName()) && + "The user's name is Alice Johnson".equals(m.getContent())); + boolean foundUserPreference = savedMemories.stream() + .anyMatch(m -> "UserPreference".equals(m.getName()) && + "User prefers email notifications over SMS".equals(m.getContent())); + + assertTrue(foundUserName); + assertTrue(foundUserPreference); + + // Verify agent name was set correctly + assertTrue(savedMemories.stream().allMatch(m -> "test-agent".equals(m.getAgentName()))); + } + + @Test + void testJsonSerializationAndStorage() throws Exception { + // Test complete JSON serialization/deserialization cycle + GeneratedMemoryUnit memory = new GeneratedMemoryUnit( + MemoryScope.AGENT, + "proc-context", + MemoryType.PROCEDURAL, + "TaskProcess", + "When handling user queries, first check existing memories for context", + List.of("process", "query", "context"), + 9 + ); + + AgentMemoryOutput original = AgentMemoryOutput.builder() + .generatedMemory(List.of(memory)) + .build(); + + // Serialize to JSON + String json = objectMapper.writeValueAsString(original); + + // Deserialize back + AgentMemoryOutput deserialized = objectMapper.readValue(json, AgentMemoryOutput.class); + + // Verify they're equal + assertEquals(original, deserialized); + + // Convert and save to storage + GeneratedMemoryUnit deserializedMemory = deserialized.getGeneratedMemory().get(0); + AgentMemory agentMemory = convertToAgentMemory(deserializedMemory, "test-agent"); + memoryStorage.save(agentMemory); + + // Verify saved in memory storage + List savedMemories = memoryStorage.findMemories( + "proc-context", MemoryScope.AGENT, null, null, null, 0, 10 + ); + + assertEquals(1, savedMemories.size()); + AgentMemory savedMemory = savedMemories.get(0); + assertEquals("TaskProcess", savedMemory.getName()); + assertEquals("When handling user queries, first check existing memories for context", + savedMemory.getContent()); + assertEquals(MemoryType.PROCEDURAL, savedMemory.getMemoryType()); + assertEquals(9, savedMemory.getReusabilityScore()); + } + + @Test + void testMultipleMemoryScopes() throws Exception { + // Test with memories from different scopes + GeneratedMemoryUnit entityMemory = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "user456", + MemoryType.SEMANTIC, + "UserJob", + "User works as a data scientist", + List.of("profession", "career"), + 8 + ); + + GeneratedMemoryUnit agentMemory = new GeneratedMemoryUnit( + MemoryScope.AGENT, + "agent-learning", + MemoryType.PROCEDURAL, + "DataScienceHelp", + "When user asks about data science, provide practical examples", + List.of("assistance", "examples"), + 7 + ); + + AgentMemoryOutput memoryOutput = AgentMemoryOutput.builder() + .generatedMemory(Arrays.asList(entityMemory, agentMemory)) + .build(); + + // Save to storage + for (GeneratedMemoryUnit generatedMemory : memoryOutput.getGeneratedMemory()) { + AgentMemory converted = convertToAgentMemory(generatedMemory, "test-agent"); + memoryStorage.save(converted); + } + + // Verify entity memories + List entityMemories = memoryStorage.findMemories( + "user456", MemoryScope.ENTITY, null, null, null, 0, 10 + ); + assertEquals(1, entityMemories.size()); + assertEquals("UserJob", entityMemories.get(0).getName()); + + // Verify agent memories + List agentMemories = memoryStorage.findMemories( + "agent-learning", MemoryScope.AGENT, null, null, null, 0, 10 + ); + assertEquals(1, agentMemories.size()); + assertEquals("DataScienceHelp", agentMemories.get(0).getName()); + } + + @Test + void testComplexMemoryStructureIntegration() throws Exception { + // Test the custom builder with complex memory structures + List complexMemories = Arrays.asList( + new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "company123", + MemoryType.SEMANTIC, + "CompanyInfo", + "TechCorp is a software company specializing in AI solutions", + Arrays.asList("company", "business", "ai", "technology"), + 9 + ), + new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "company123", + MemoryType.EPISODIC, + "RecentMeeting", + "Had a productive meeting about Q4 goals last Tuesday", + Arrays.asList("meeting", "goals", "quarterly"), + 6 + ), + new GeneratedMemoryUnit( + MemoryScope.AGENT, + "business-context", + MemoryType.PROCEDURAL, + "BusinessCommunication", + "When discussing business topics, use professional language and focus on ROI", + Arrays.asList("communication", "business", "professional"), + 8 + ) + ); + + AgentMemoryOutput complexOutput = AgentMemoryOutput.builder() + .generatedMemory(complexMemories) + .build(); + + // Verify serialization/deserialization works + String json = objectMapper.writeValueAsString(complexOutput); + AgentMemoryOutput deserialized = objectMapper.readValue(json, AgentMemoryOutput.class); + assertEquals(complexOutput, deserialized); + + // Save to storage + for (GeneratedMemoryUnit generatedMemory : deserialized.getGeneratedMemory()) { + AgentMemory converted = convertToAgentMemory(generatedMemory, "test-agent"); + memoryStorage.save(converted); + } + + // Verify all memories were saved correctly + List companyMemories = memoryStorage.findMemories( + "company123", MemoryScope.ENTITY, null, null, null, 0, 10 + ); + assertEquals(2, companyMemories.size()); + + List agentMemories = memoryStorage.findMemories( + "business-context", MemoryScope.AGENT, null, null, null, 0, 10 + ); + assertEquals(1, agentMemories.size()); + + // Verify all topics are preserved + AgentMemory companyInfo = companyMemories.stream() + .filter(m -> "CompanyInfo".equals(m.getName())) + .findFirst() + .orElseThrow(); + assertTrue(companyInfo.getTopics().contains("ai")); + assertTrue(companyInfo.getTopics().contains("technology")); + assertEquals(4, companyInfo.getTopics().size()); + } + + @Test + void testEmptyAndNullMemoryLists() throws Exception { + // Test with empty memory list + AgentMemoryOutput emptyOutput = AgentMemoryOutput.builder() + .generatedMemory(List.of()) + .build(); + + assertNotNull(emptyOutput); + assertTrue(emptyOutput.getGeneratedMemory().isEmpty()); + + // Test with null memory list + AgentMemoryOutput nullOutput = AgentMemoryOutput.builder() + .generatedMemory(null) + .build(); + + assertNotNull(nullOutput); + assertNull(nullOutput.getGeneratedMemory()); + + // Verify storage remains empty + List allMemories = memoryStorage.findMemories( + null, null, null, null, null, 0, 100 + ); + assertTrue(allMemories.isEmpty()); + } + + @Test + void testBuilderCustomClassNameWithStorage() { + // Test that the custom builder class name works correctly with storage operations + AgentMemoryOutput.AgentMemoryOutputBuilder builder = AgentMemoryOutput.builder(); + assertEquals("AgentMemoryOutputBuilder", builder.getClass().getSimpleName()); + + GeneratedMemoryUnit testMemory = new GeneratedMemoryUnit( + MemoryScope.ENTITY, + "builder-test", + MemoryType.SEMANTIC, + "BuilderTest", + "Testing custom builder integration", + List.of("test", "builder"), + 7 + ); + + AgentMemoryOutput output = builder.generatedMemory(List.of(testMemory)).build(); + + // Save to storage + AgentMemory converted = convertToAgentMemory(testMemory, "builder-test-agent"); + memoryStorage.save(converted); + + // Verify it was saved correctly + List saved = memoryStorage.findMemories( + "builder-test", MemoryScope.ENTITY, null, null, null, 0, 10 + ); + assertEquals(1, saved.size()); + assertEquals("BuilderTest", saved.get(0).getName()); + } + + /** + * Helper method to convert GeneratedMemoryUnit to AgentMemory + */ + private AgentMemory convertToAgentMemory(GeneratedMemoryUnit generatedMemory, String agentName) { + return AgentMemory.builder() + .agentName(agentName) + .scope(generatedMemory.getScope()) + .scopeId(generatedMemory.getScopeId()) + .memoryType(generatedMemory.getType()) + .name(generatedMemory.getName()) + .content(generatedMemory.getContent()) + .topics(generatedMemory.getTopics()) + .reusabilityScore(generatedMemory.getReusabilityScore()) + .build(); + } +} diff --git a/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorageTest.java b/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorageTest.java new file mode 100644 index 0000000..755fad7 --- /dev/null +++ b/sentinel-ai-storage-es/src/test/java/com/phonepe/sentinelai/storage/memory/InMemoryMemoryStorageTest.java @@ -0,0 +1,469 @@ +package com.phonepe.sentinelai.storage.memory; + +import com.phonepe.sentinelai.agentmemory.AgentMemory; +import com.phonepe.sentinelai.agentmemory.MemoryScope; +import com.phonepe.sentinelai.agentmemory.MemoryType; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.time.LocalDateTime; +import java.util.*; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.stream.IntStream; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Tests for {@link InMemoryMemoryStorage} + */ +@Slf4j +class InMemoryMemoryStorageTest { + + private InMemoryMemoryStorage memoryStorage; + + @BeforeEach + void setUp() { + memoryStorage = new InMemoryMemoryStorage(); + } + + @Test + void testSaveMemory() { + // Test saving a memory + AgentMemory memory = createTestMemory( + "test-agent", + MemoryScope.ENTITY, + "user123", + MemoryType.SEMANTIC, + "UserName", + "User's name is John Doe", + List.of("personal", "name"), + 8 + ); + + Optional saved = memoryStorage.save(memory); + + assertTrue(saved.isPresent()); + assertEquals(memory, saved.get()); + } + + @Test + void testFindMemoriesWithValidScopeAndScopeId() { + // Save test memories + AgentMemory memory1 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + AgentMemory memory2 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserAge", "25 years old", List.of("personal"), 7 + ); + AgentMemory memory3 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user456", MemoryType.SEMANTIC, + "UserName", "Jane Smith", List.of("personal"), 6 + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + memoryStorage.save(memory3); + + // Find memories for user123 + List memories = memoryStorage.findMemories( + "user123", MemoryScope.ENTITY, Set.of(MemoryType.SEMANTIC), + List.of(), null, 0, 10 + ); + + assertEquals(2, memories.size()); + assertTrue(memories.contains(memory1)); + assertTrue(memories.contains(memory2)); + assertFalse(memories.contains(memory3)); + } + + @Test + void testFindMemoriesWithNullScopeAndScopeId() { + // Save test memories with different scopes + AgentMemory memory1 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + AgentMemory memory2 = createTestMemory( + "agent2", MemoryScope.AGENT, "proc-scope", MemoryType.PROCEDURAL, + "Process", "How to handle requests", List.of("procedure"), 7 + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + + // Find all memories when scope and scopeId are null + List allMemories = memoryStorage.findMemories( + null, null, Set.of(), List.of(), null, 0, 10 + ); + + assertEquals(2, allMemories.size()); + assertTrue(allMemories.contains(memory1)); + assertTrue(allMemories.contains(memory2)); + } + + @Test + void testFindMemoriesWithNullScope() { + // Save test memories + AgentMemory memory = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + memoryStorage.save(memory); + + // Find memories with null scope but valid scopeId + List memories = memoryStorage.findMemories( + "user123", null, Set.of(), List.of(), null, 0, 10 + ); + + // Should return all memories since scope is null + assertEquals(1, memories.size()); + assertTrue(memories.contains(memory)); + } + + @Test + void testFindMemoriesWithNullScopeId() { + // Save test memories + AgentMemory memory = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + memoryStorage.save(memory); + + // Find memories with valid scope but null scopeId + List memories = memoryStorage.findMemories( + null, MemoryScope.ENTITY, Set.of(), List.of(), null, 0, 10 + ); + + // Should return all memories since scopeId is null + assertEquals(1, memories.size()); + assertTrue(memories.contains(memory)); + } + + @Test + void testFindMemoriesEmptyResult() { + // Save a memory + AgentMemory memory = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + memoryStorage.save(memory); + + // Try to find memories for non-existent scope/scopeId combination + List memories = memoryStorage.findMemories( + "nonexistent", MemoryScope.AGENT, Set.of(), List.of(), null, 0, 10 + ); + + assertTrue(memories.isEmpty()); + } + + @Test + void testSaveMultipleMemoriesInSameScope() { + // Save multiple memories in the same scope + String scopeId = "user123"; + MemoryScope scope = MemoryScope.ENTITY; + + AgentMemory memory1 = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + AgentMemory memory2 = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "UserAge", "25", List.of("personal"), 7 + ); + AgentMemory memory3 = createTestMemory( + "agent1", scope, scopeId, MemoryType.EPISODIC, + "LastLogin", "Logged in yesterday", List.of("activity"), 6 + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + memoryStorage.save(memory3); + + List memories = memoryStorage.findMemories( + scopeId, scope, Set.of(), List.of(), null, 0, 10 + ); + + assertEquals(3, memories.size()); + assertTrue(memories.contains(memory1)); + assertTrue(memories.contains(memory2)); + assertTrue(memories.contains(memory3)); + } + + @Test + void testConcurrentAccess() throws Exception { + // Test concurrent access to the memory storage + int numThreads = 10; + int memoriesPerThread = 20; + ExecutorService executor = Executors.newFixedThreadPool(numThreads); + CountDownLatch latch = new CountDownLatch(numThreads); + + List> futures = new ArrayList<>(); + + for (int i = 0; i < numThreads; i++) { + final int threadId = i; + CompletableFuture future = CompletableFuture.runAsync(() -> { + try { + // Save memories + for (int j = 0; j < memoriesPerThread; j++) { + AgentMemory memory = createTestMemory( + "agent" + threadId, + MemoryScope.ENTITY, + "user" + threadId, + MemoryType.SEMANTIC, + "Memory" + j, + "Content " + j, + List.of("test"), + 5 + ); + memoryStorage.save(memory); + } + + // Read memories + List memories = memoryStorage.findMemories( + "user" + threadId, MemoryScope.ENTITY, + Set.of(), List.of(), null, 0, 100 + ); + assertEquals(memoriesPerThread, memories.size()); + } finally { + latch.countDown(); + } + }, executor); + futures.add(future); + } + + // Wait for all threads to complete + latch.await(); + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); + + // Verify total memories saved + List allMemories = memoryStorage.findMemories( + null, null, Set.of(), List.of(), null, 0, 1000 + ); + assertEquals(numThreads * memoriesPerThread, allMemories.size()); + + executor.shutdown(); + } + + @Test + void testMemoryOverwriteInSameScope() { + // Test that memories accumulate in the same scope (no overwrite) + String scopeId = "user123"; + MemoryScope scope = MemoryScope.ENTITY; + + AgentMemory memory1 = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "UserName", "John Doe", List.of("personal"), 8 + ); + + AgentMemory memory2 = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "UserName", "John Smith", List.of("personal"), 9 // Different content, same name + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + + List memories = memoryStorage.findMemories( + scopeId, scope, Set.of(), List.of(), null, 0, 10 + ); + + // Both memories should be present (no overwrite) + assertEquals(2, memories.size()); + assertTrue(memories.contains(memory1)); + assertTrue(memories.contains(memory2)); + } + + @Test + void testDifferentMemoryTypes() { + // Test with different memory types + String scopeId = "test-scope"; + MemoryScope scope = MemoryScope.AGENT; + + AgentMemory semanticMemory = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "Fact", "Important fact", List.of("facts"), 8 + ); + + AgentMemory proceduralMemory = createTestMemory( + "agent1", scope, scopeId, MemoryType.PROCEDURAL, + "Process", "How to do something", List.of("procedure"), 7 + ); + + AgentMemory episodicMemory = createTestMemory( + "agent1", scope, scopeId, MemoryType.EPISODIC, + "Event", "What happened", List.of("events"), 6 + ); + + memoryStorage.save(semanticMemory); + memoryStorage.save(proceduralMemory); + memoryStorage.save(episodicMemory); + + List allMemories = memoryStorage.findMemories( + scopeId, scope, Set.of(), List.of(), null, 0, 10 + ); + + assertEquals(3, allMemories.size()); + assertTrue(allMemories.contains(semanticMemory)); + assertTrue(allMemories.contains(proceduralMemory)); + assertTrue(allMemories.contains(episodicMemory)); + } + + @Test + void testMemoryParametersAreIgnored() { + // Test that memory type, topics, query, minReusabilityScore, and count parameters are ignored + // in the current implementation + + AgentMemory memory1 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.SEMANTIC, + "UserName", "John", List.of("personal"), 8 + ); + AgentMemory memory2 = createTestMemory( + "agent1", MemoryScope.ENTITY, "user123", MemoryType.PROCEDURAL, + "Process", "How to", List.of("procedure"), 3 // Low reusability score + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + + // Find with specific memory types - should still return all memories for the scope + List semanticOnly = memoryStorage.findMemories( + "user123", MemoryScope.ENTITY, Set.of(MemoryType.SEMANTIC), + List.of(), null, 5, 1 // min score 5, count 1 + ); + + // Current implementation returns all memories regardless of filters + assertEquals(2, semanticOnly.size()); + } + + @Test + void testEmptyMemoryStorage() { + // Test operations on empty storage + List memories = memoryStorage.findMemories( + "any", MemoryScope.ENTITY, Set.of(), List.of(), null, 0, 10 + ); + assertTrue(memories.isEmpty()); + + List allMemories = memoryStorage.findMemories( + null, null, Set.of(), List.of(), null, 0, 10 + ); + assertTrue(allMemories.isEmpty()); + } + + @Test + void testLargeNumberOfMemories() { + // Test with a large number of memories + String scopeId = "load-test"; + MemoryScope scope = MemoryScope.ENTITY; + int numMemories = 1000; + + List savedMemories = new ArrayList<>(); + for (int i = 0; i < numMemories; i++) { + AgentMemory memory = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "Memory" + i, "Content " + i, List.of("test"), i % 10 + ); + memoryStorage.save(memory); + savedMemories.add(memory); + } + + List retrievedMemories = memoryStorage.findMemories( + scopeId, scope, Set.of(), List.of(), null, 0, numMemories + 100 + ); + + assertEquals(numMemories, retrievedMemories.size()); + assertTrue(retrievedMemories.containsAll(savedMemories)); + } + + @Test + void testThreadSafetyWithConcurrentHashMap() { + // Verify that ConcurrentHashMap behavior is maintained + ExecutorService executor = Executors.newFixedThreadPool(5); + + List> futures = IntStream.range(0, 100) + .mapToObj(i -> CompletableFuture.runAsync(() -> { + AgentMemory memory = createTestMemory( + "agent" + (i % 5), + MemoryScope.ENTITY, + "user" + (i % 10), + MemoryType.SEMANTIC, + "Memory" + i, + "Content " + i, + List.of("concurrent"), + 5 + ); + memoryStorage.save(memory); + }, executor)) + .toList(); + + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); + + // Verify all memories were saved + List allMemories = memoryStorage.findMemories( + null, null, Set.of(), List.of(), null, 0, 1000 + ); + assertEquals(100, allMemories.size()); + + executor.shutdown(); + } + + @Test + void testKeyRecordEquality() { + // Test that the internal Key record works correctly + String scopeId = "test"; + MemoryScope scope = MemoryScope.ENTITY; + + AgentMemory memory1 = createTestMemory( + "agent1", scope, scopeId, MemoryType.SEMANTIC, + "Test1", "Content1", List.of("test"), 5 + ); + AgentMemory memory2 = createTestMemory( + "agent2", scope, scopeId, MemoryType.PROCEDURAL, + "Test2", "Content2", List.of("test"), 6 + ); + + memoryStorage.save(memory1); + memoryStorage.save(memory2); + + // Both memories should be in the same scope + List memories = memoryStorage.findMemories( + scopeId, scope, Set.of(), List.of(), null, 0, 10 + ); + + assertEquals(2, memories.size()); + assertTrue(memories.contains(memory1)); + assertTrue(memories.contains(memory2)); + } + + /** + * Helper method to create test AgentMemory instances + */ + private AgentMemory createTestMemory( + String agentName, + MemoryScope scope, + String scopeId, + MemoryType memoryType, + String name, + String content, + List topics, + int reusabilityScore) { + + return AgentMemory.builder() + .agentName(agentName) + .scope(scope) + .scopeId(scopeId) + .memoryType(memoryType) + .name(name) + .content(content) + .topics(topics) + .reusabilityScore(reusabilityScore) + .createdAt(LocalDateTime.now()) + .updatedAt(LocalDateTime.now()) + .build(); + } +}