Chat Memory Management

// Package: io.quarkiverse.langchain4j.runtime.aiservice
/**
 * Interface for seeding chat memory with example messages.
 * Used to provide few-shot learning examples or conversation context.
 * 
 * Implementations are CDI beans invoked once per memory ID
 * when ChatMemory is first created.
 * 
 * Thread Safety: Implementations must be thread-safe.
 * 
 * Lifecycle: Called synchronously during memory creation,
 * before first user message is added.
 */
public interface ChatMemorySeeder {

    /**
     * Seed chat memory with example messages.
     * 
     * Messages are added to memory in the order returned.
     * Typical pattern: UserMessage, AiMessage, UserMessage, AiMessage, ...
     * 
     * System messages can also be included but are usually
     * handled via @SystemMessage annotation instead.
     * 
     * Return empty list to skip seeding for certain contexts.
     * Throw exception to fail memory creation.
     *
     * @param context Context containing method name and other metadata
     * @return List of chat messages to seed into memory
     * @throws RuntimeException to fail memory creation
     */
    List<ChatMessage> seed(Context context);

    /**
     * Context passed to seed method.
     * Provides information about why memory is being created.
     * 
     * @param methodName Name of the AI service method being invoked
     */
    record Context(String methodName) {}
}

import io.quarkiverse.langchain4j.runtime.aiservice.ChatMemorySeeder;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.AiMessage;
import jakarta.enterprise.context.ApplicationScoped;

@ApplicationScoped
public class ExampleSeeder implements ChatMemorySeeder {

    @Override
    public List<ChatMessage> seed(Context context) {
        if ("reviewCode".equals(context.methodName())) {
            return List.of(
                UserMessage.from("def add(a, b): return a + b"),
                AiMessage.from("Good: clear function name and simple implementation."),
                UserMessage.from("def x(a,b):return a+b"),
                AiMessage.from("Issues: unclear name 'x', missing spaces per PEP 8.")
            );
        }
        return List.of();
    }
}

@ApplicationScoped
public class MethodAwareSeeder implements ChatMemorySeeder {

    @Inject
    ExampleRepository exampleRepo;

    @Override
    public List<ChatMessage> seed(Context context) {
        // Load examples from database based on method
        return switch (context.methodName()) {
            case "translateText" -> exampleRepo.findTranslationExamples();
            case "summarizeDocument" -> exampleRepo.findSummarizationExamples();
            case "answerQuestion" -> exampleRepo.findQuestionAnswerExamples();
            default -> List.of();
        };
    }
}

// Package: io.quarkiverse.langchain4j
/**
 * Utility for manual chat memory removal.
 * Use when you need to clear conversation history programmatically.
 * 
 * Static methods for convenience - no instance needed.
 * 
 * Use cases:
 * - User logout
 * - Privacy compliance (GDPR right to be forgotten)
 * - Session cleanup
 * - Testing cleanup
 * - Force memory reset
 */
public final class ChatMemoryRemover {

    /**
     * Remove chat memory by ID.
     * 
     * Calls removeChatMemoryIds() on the AI service's QuarkusAiServiceContext.
     * If ChatMemoryProvider implements ChatMemoryRemovable, also calls remove().
     * 
     * Behavior:
     * - Removes from in-memory cache
     * - Optionally removes from persistent storage
     * - Next access creates fresh memory (with seeding)
     * - No-op if memory ID doesn't exist
     *
     * @param aiService The AI service instance (must be @RegisterAiService)
     * @param memoryId The memory ID to remove
     * @throws IllegalArgumentException if aiService is not a registered AI service
     */
    public static void remove(Object aiService, Object memoryId);

    /**
     * Remove multiple chat memory IDs.
     * 
     * More efficient than calling remove() multiple times
     * as it can batch operations.
     * 
     * Behavior:
     * - All IDs processed in single batch
     * - Failures for individual IDs don't stop others
     * - No exception if some IDs don't exist
     *
     * @param aiService The AI service instance
     * @param memoryIds List of memory IDs to remove
     * @throws IllegalArgumentException if aiService is not a registered AI service
     */
    public static void remove(Object aiService, List<Object> memoryIds);
}

import io.quarkiverse.langchain4j.ChatMemoryRemover;
import jakarta.inject.Inject;

public class ConversationManager {

    @Inject
    AssistantService assistant;

    public void clearUserConversation(String userId) {
        // Remove single memory ID
        ChatMemoryRemover.remove(assistant, userId);
    }

    public void clearMultipleConversations(List<String> userIds) {
        // Remove multiple memory IDs efficiently
        ChatMemoryRemover.remove(assistant, userIds);
    }
    
    public void clearAllConversations() {
        // Get all active user IDs and remove
        List<String> allUserIds = userRepository.findAllActiveUserIds();
        ChatMemoryRemover.remove(assistant, allUserIds);
    }
}

@ApplicationScoped
public class MemoryCleanupService {

    @Inject
    Assistant assistant;

    @Inject
    SessionRepository sessionRepo;

    @Scheduled(every = "1h")
    public void cleanupExpiredSessions() {
        // Find sessions older than 24 hours
        List<String> expiredSessions = sessionRepo.findExpiredSessions(
            Instant.now().minus(24, ChronoUnit.HOURS)
        );

        if (!expiredSessions.isEmpty()) {
            ChatMemoryRemover.remove(assistant, expiredSessions);
            logger.info("Cleaned up {} expired conversation memories", expiredSessions.size());
        }
    }
}

// Package: io.quarkiverse.langchain4j
/**
 * Interface for chat memory that supports removal of conversation history.
 * Implement this to enable programmatic memory cleanup.
 * 
 * ChatMemoryProvider implementations that support persistence
 * should also implement this interface to handle removal.
 * 
 * If not implemented, ChatMemoryRemover will only evict from cache
 * but won't remove from persistent storage.
 */
public interface ChatMemoryRemovable {

    /**
     * Remove chat memories by IDs.
     * 
     * Called by ChatMemoryRemover.remove().
     * Should remove from persistent storage if applicable.
     * 
     * Implementation guidelines:
     * - Handle batch removal efficiently
     * - Don't throw exception for non-existent IDs
     * - Be idempotent (safe to call multiple times)
     * - Be thread-safe
     * - Log failures but don't propagate for individual IDs
     *
     * @param ids Variable number of memory IDs to remove
     */
    void remove(Object... ids);
}

import io.quarkiverse.langchain4j.ChatMemoryRemovable;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.ChatMemoryProvider;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@ApplicationScoped
public class PostgresChatMemoryProvider 
    implements ChatMemoryProvider, ChatMemoryRemovable {

    @Inject
    DataSource dataSource;

    private final Map<Object, ChatMemory> cache = new ConcurrentHashMap<>();

    @Override
    public ChatMemory get(Object memoryId) {
        return cache.computeIfAbsent(memoryId, id -> {
            // Load from database or create new
            return loadOrCreateMemory(id);
        });
    }

    @Override
    public void remove(Object... ids) {
        try (Connection conn = dataSource.getConnection()) {
            String sql = "DELETE FROM chat_memory WHERE memory_id = ?";
            try (PreparedStatement stmt = conn.prepareStatement(sql)) {
                for (Object id : ids) {
                    // Remove from cache
                    cache.remove(id);
                    
                    // Remove from database
                    stmt.setString(1, id.toString());
                    stmt.addBatch();
                }
                stmt.executeBatch();
            }
        } catch (SQLException e) {
            logger.error("Failed to remove chat memories", e);
            // Don't propagate exception
        }
    }

    private ChatMemory loadOrCreateMemory(Object memoryId) {
        // Implementation to load from database
        return new MessageWindowChatMemory(100);
    }
}

import io.quarkiverse.langchain4j.ChatMemoryRemovable;
import dev.langchain4j.memory.ChatMemoryProvider;
import redis.clients.jedis.JedisPool;

@ApplicationScoped
public class RedisChatMemoryProvider 
    implements ChatMemoryProvider, ChatMemoryRemovable {

    @Inject
    JedisPool jedisPool;

    @Override
    public ChatMemory get(Object memoryId) {
        // Load from Redis
        return new RedisChatMemory(jedisPool, memoryId.toString());
    }

    @Override
    public void remove(Object... ids) {
        try (Jedis jedis = jedisPool.getResource()) {
            String[] keys = Arrays.stream(ids)
                .map(id -> "chat:memory:" + id)
                .toArray(String[]::new);
            jedis.del(keys);
        } catch (Exception e) {
            logger.error("Failed to remove chat memories from Redis", e);
        }
    }
}

// Package: io.quarkiverse.langchain4j.runtime.aiservice
/**
 * Get chat memory by ID.
 * Creates new memory if it doesn't exist.
 * 
 * Process:
 * 1. Check cache for existing memory
 * 2. If not found, call ChatMemoryProvider.get()
 * 3. If seeder configured, call seeder.seed()
 * 4. Add seed messages to memory
 * 5. Store in cache
 * 6. Return memory
 * 
 * Thread Safety: Synchronized to prevent duplicate creation.
 * Performance: O(1) for cached memories, O(n) for new creation.
 *
 * @param id The memory ID (typically user ID or session ID)
 * @return The chat memory for this ID (never null)
 * @throws RuntimeException if memory creation fails
 */
public ChatMemory getChatMemory(Object id);

/**
 * Evict chat memory from cache.
 * Memory will be reloaded on next access.
 * 
 * Use case: Force memory reload after external changes.
 * Does NOT call remove() on the memory provider.
 * Does NOT call ChatMemoryRemovable.remove().
 * 
 * Behavior:
 * - Removes from cache only
 * - Next access creates/loads fresh memory
 * - Seeder will run again for new memory
 * - Persistent storage unchanged
 * 
 * Thread Safety: Safe to call concurrently.
 *
 * @param id The memory ID to evict
 */
public void evictChatMemory(Object id);

/**
 * Remove multiple chat memory IDs.
 * Permanently deletes conversation history.
 * 
 * Process:
 * 1. Evict from cache
 * 2. If ChatMemoryProvider implements ChatMemoryRemovable:
 *    - Call remove() to delete from storage
 * 3. Otherwise, only eviction occurs
 * 
 * Use case: User logout, session cleanup, privacy compliance.
 * 
 * Thread Safety: Safe to call concurrently.
 * Performance: O(n) where n is number of IDs.
 *
 * @param ids Memory IDs to remove
 */
public void removeChatMemoryIds(Object... ids);

// Package: io.quarkiverse.langchain4j.runtime.aiservice
/**
 * Close context and clean up resources.
 * Called during application shutdown.
 * 
 * Behavior:
 * - Clears all cached memories
 * - Closes model instances
 * - Releases resources
 * - Does NOT remove memories from storage
 * 
 * After close(), context is unusable.
 */
@Override
public void close();

import io.quarkiverse.langchain4j.RegisterAiService;

// Use CDI bean for memory provider (default)
@RegisterAiService(
    chatMemoryProviderSupplier = RegisterAiService.BeanChatMemoryProviderSupplier.class
)
public interface WithMemory {
    String chat(String message);
}

// Disable chat memory - stateless service
@RegisterAiService(
    chatMemoryProviderSupplier = RegisterAiService.NoChatMemoryProviderSupplier.class
)
public interface WithoutMemory {
    String chat(String message);
}

// Custom memory provider
@RegisterAiService(
    chatMemoryProviderSupplier = CustomMemoryProviderSupplier.class
)
public interface CustomMemory {
    String chat(String message);
}

import java.util.function.Supplier;
import dev.langchain4j.memory.ChatMemoryProvider;

public class CustomMemoryProviderSupplier implements Supplier<ChatMemoryProvider> {
    
    @Override
    public ChatMemoryProvider get() {
        // Return custom provider instance
        return new MyCustomChatMemoryProvider();
    }
}

import io.quarkiverse.langchain4j.RegisterAiService;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;

@RegisterAiService
public interface Assistant {

    // Memory ID automatically extracted from @MemoryId parameter
    String chat(@MemoryId String userId, @UserMessage String message);

    // Default memory ID used if no @MemoryId parameter
    // All invocations share the same memory
    String chat(@UserMessage String message);
    
    // Memory ID can be any object type
    String chat(@MemoryId UUID sessionId, String message);
    
    // Multiple parameters, but only one @MemoryId
    String chat(@MemoryId String userId, String context, String message);
}

// Good: Same user ID = same memory
assistant.chat("user123", "Hello");
assistant.chat("user123", "How are you?");  // Continues conversation

// Bad: Different memory IDs = separate conversations
assistant.chat("user123", "Hello");
assistant.chat("user_123", "How are you?");  // New conversation!

import dev.langchain4j.memory.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import jakarta.enterprise.context.ApplicationScoped;

@ApplicationScoped
public class SimpleChatMemoryProvider implements ChatMemoryProvider {

    private static final int MAX_MESSAGES = 100;

    @Override
    public ChatMemory get(Object memoryId) {
        // Create new memory with message window
        return MessageWindowChatMemory.withMaxMessages(MAX_MESSAGES);
    }
}

import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

@ApplicationScoped
public class PersistentChatMemoryProvider implements ChatMemoryProvider {

    private final Map<Object, ChatMemory> memories = new ConcurrentHashMap<>();

    @Override
    public ChatMemory get(Object memoryId) {
        return memories.computeIfAbsent(memoryId, 
            id -> MessageWindowChatMemory.withMaxMessages(100));
    }
}

@ApplicationScoped
public class DatabaseChatMemoryProvider 
    implements ChatMemoryProvider, ChatMemoryRemovable {

    @Inject
    ChatMemoryRepository repository;

    private final Map<Object, ChatMemory> cache = new ConcurrentHashMap<>();

    @Override
    public ChatMemory get(Object memoryId) {
        return cache.computeIfAbsent(memoryId, this::loadFromDatabase);
    }

    private ChatMemory loadFromDatabase(Object memoryId) {
        List<ChatMessage> messages = repository.findByMemoryId(memoryId.toString());
        ChatMemory memory = MessageWindowChatMemory.withMaxMessages(100);
        messages.forEach(memory::add);
        return memory;
    }

    @Override
    public void remove(Object... ids) {
        for (Object id : ids) {
            cache.remove(id);
            repository.deleteByMemoryId(id.toString());
        }
    }
}

@ApplicationScoped
public class RedisChatMemoryProvider 
    implements ChatMemoryProvider, ChatMemoryRemovable {

    @Inject
    @ConfigProperty(name = "redis.host")
    String redisHost;

    @Inject
    @ConfigProperty(name = "redis.port")
    int redisPort;

    private JedisPool jedisPool;

    @PostConstruct
    void init() {
        jedisPool = new JedisPool(redisHost, redisPort);
    }

    @Override
    public ChatMemory get(Object memoryId) {
        String key = "chat:memory:" + memoryId;
        
        try (Jedis jedis = jedisPool.getResource()) {
            List<String> serializedMessages = jedis.lrange(key, 0, -1);
            ChatMemory memory = MessageWindowChatMemory.withMaxMessages(100);
            
            for (String serialized : serializedMessages) {
                ChatMessage message = deserializeMessage(serialized);
                memory.add(message);
            }
            
            return new RedisBackedChatMemory(jedisPool, key, memory);
        }
    }

    @Override
    public void remove(Object... ids) {
        try (Jedis jedis = jedisPool.getResource()) {
            for (Object id : ids) {
                String key = "chat:memory:" + id;
                jedis.del(key);
            }
        }
    }

    private ChatMessage deserializeMessage(String serialized) {
        // JSON deserialization
        return objectMapper.readValue(serialized, ChatMessage.class);
    }

    @PreDestroy
    void cleanup() {
        if (jedisPool != null) {
            jedisPool.close();
        }
    }
}

1. AI service method called with @MemoryId parameter
2. Context.getChatMemory(memoryId) called
3. Check cache for existing memory
4. If not in cache:
   a. Call ChatMemoryProvider.get(memoryId)
   b. If ChatMemorySeeder configured:
      - Call seeder.seed(context)
      - Add seed messages to memory
   c. Store in cache
5. Add current user message to memory
6. Return memory

First call:  Create + Seed + Use
Second call: Use (from cache)
Third call:  Use (from cache)
After evict: Create + Seed + Use (seeding again)
After remove: Create + Seed + Use (fresh memory)

// Message window: Keep last N messages
ChatMemory memory = MessageWindowChatMemory.withMaxMessages(100);

// Token window: Keep messages up to token limit
ChatMemory memory = MessageWindowChatMemory.withMaxTokens(4000, tokenizer);

@ApplicationScoped
public class TokenLimitedMemoryProvider implements ChatMemoryProvider {

    @Inject
    Tokenizer tokenizer;  // Inject from LangChain4j

    @Override
    public ChatMemory get(Object memoryId) {
        // Limit to 4000 tokens (e.g., for GPT-3.5 with 4096 token limit)
        return MessageWindowChatMemory.builder()
            .id(memoryId)
            .maxTokens(4000, tokenizer)
            .build();
    }
}

@RegisterAiService
public interface SmartAssistant {

    // Persistent memory (user ID)
    String chat(@MemoryId String userId, String message);
    
    // Temporary memory (session ID, cleaned up on logout)
    String temporaryChat(@MemoryId UUID sessionId, String message);
}

public record HierarchicalMemoryId(String tenantId, String userId, String sessionId) {
    
    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (!(o instanceof HierarchicalMemoryId that)) return false;
        return Objects.equals(tenantId, that.tenantId) &&
               Objects.equals(userId, that.userId) &&
               Objects.equals(sessionId, that.sessionId);
    }

    @Override
    public int hashCode() {
        return Objects.hash(tenantId, userId, sessionId);
    }
}

@RegisterAiService
public interface MultiTenantAssistant {
    String chat(@MemoryId HierarchicalMemoryId id, String message);
}

@ApplicationScoped
public class ConditionalSeeder implements ChatMemorySeeder {

    @Inject
    UserPreferencesService preferences;

    @Override
    public List<ChatMessage> seed(Context context) {
        // Get current user from context (e.g., thread local, request scope)
        String userId = SecurityContext.getCurrentUserId();
        
        // Check user preferences
        if (preferences.wantsFewShotExamples(userId)) {
            return loadExamples(context.methodName());
        }
        
        return List.of();
    }

    private List<ChatMessage> loadExamples(String methodName) {
        // Load method-specific examples
        return exampleRepository.findByMethod(methodName);
    }
}

@ApplicationScoped
public class DynamicWindowMemoryProvider implements ChatMemoryProvider {

    @Inject
    SubscriptionService subscriptionService;

    @Override
    public ChatMemory get(Object memoryId) {
        // Different memory limits based on subscription tier
        String tier = subscriptionService.getUserTier(memoryId.toString());
        
        int maxMessages = switch (tier) {
            case "premium" -> 200;
            case "standard" -> 100;
            case "free" -> 50;
            default -> 20;
        };
        
        return MessageWindowChatMemory.withMaxMessages(maxMessages);
    }
}

@ApplicationScoped
public class ExpiringMemoryProvider 
    implements ChatMemoryProvider, ChatMemoryRemovable {

    private final Map<Object, TimestampedMemory> memories = new ConcurrentHashMap<>();
    private static final long EXPIRATION_MS = 3600_000; // 1 hour

    @Override
    public ChatMemory get(Object memoryId) {
        TimestampedMemory timestamped = memories.compute(memoryId, (id, existing) -> {
            long now = System.currentTimeMillis();
            
            if (existing == null || (now - existing.timestamp) > EXPIRATION_MS) {
                // Create new memory if expired or missing
                return new TimestampedMemory(
                    MessageWindowChatMemory.withMaxMessages(100),
                    now
                );
            }
            
            // Update timestamp on access
            return new TimestampedMemory(existing.memory, now);
        });
        
        return timestamped.memory;
    }

    @Override
    public void remove(Object... ids) {
        for (Object id : ids) {
            memories.remove(id);
        }
    }

    @Scheduled(every = "15m")
    void cleanupExpired() {
        long now = System.currentTimeMillis();
        memories.entrySet().removeIf(entry -> 
            (now - entry.getValue().timestamp) > EXPIRATION_MS
        );
    }

    private record TimestampedMemory(ChatMemory memory, long timestamp) {}
}

@Test
void testMemoryPersistence() {
    SimpleChatMemoryProvider provider = new SimpleChatMemoryProvider();
    
    ChatMemory memory1 = provider.get("user1");
    memory1.add(UserMessage.from("Hello"));
    
    ChatMemory memory2 = provider.get("user1");
    assertEquals(1, memory2.messages().size());
    assertEquals("Hello", memory2.messages().get(0).text());
}

@QuarkusTest
class AssistantMemoryTest {

    @Inject
    Assistant assistant;

    @Test
    void testConversationContinuity() {
        String userId = "test-user-" + UUID.randomUUID();
        
        // First message
        String response1 = assistant.chat(userId, "My name is Alice");
        
        // Second message should remember first
        String response2 = assistant.chat(userId, "What is my name?");
        assertTrue(response2.toLowerCase().contains("alice"));
    }

    @AfterEach
    void cleanup() {
        // Clean up test memories
        ChatMemoryRemover.remove(assistant, testUserIds);
    }
}