Chat Memory Repositories

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-model-chat-memory-repository-jdbc</artifactId>
</dependency>

spring:
  datasource:
    url: jdbc:postgresql://localhost:5432/docsassistant
    username: ${DOCSASSISTANT_DB_USER}
    password: ${DOCSASSISTANT_DB_PASSWORD}
  ai:
    chat:
      memory:
        repository:
          jdbc:
            initialize-schema: always   # development only; use "never" with a migration tool in production
docsassistant:
  memory:
    max-messages: 20

CREATE TABLE IF NOT EXISTS SPRING_AI_CHAT_MEMORY (
    conversation_id VARCHAR(36) NOT NULL,
    content TEXT NOT NULL,
    type VARCHAR(10) NOT NULL CHECK (type IN ('USER', 'ASSISTANT', 'SYSTEM', 'TOOL')),
    "timestamp" TIMESTAMP NOT NULL,
    sequence_id BIGINT NOT NULL
    );

CREATE INDEX IF NOT EXISTS SPRING_AI_CHAT_MEMORY_CONVERSATION_ID_TIMESTAMP_IDX
ON SPRING_AI_CHAT_MEMORY(conversation_id, "timestamp");

CREATE INDEX IF NOT EXISTS SPRING_AI_CHAT_MEMORY_CONVERSATION_ID_SEQUENCE_ID_IDX
ON SPRING_AI_CHAT_MEMORY(conversation_id, sequence_id);

package com.example.docsassistant;

import javax.sql.DataSource;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepositoryDialect;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.transaction.PlatformTransactionManager;

/** Only needed when you do not want the auto-configured repository, for example to share a transaction manager. */
@Configuration(proxyBeanMethods = false)
class JdbcMemoryConfig {

    @Bean
    JdbcChatMemoryRepository chatMemoryRepository(DataSource dataSource, PlatformTransactionManager txManager) {
        return JdbcChatMemoryRepository.builder()
                .jdbcTemplate(new JdbcTemplate(dataSource))
                .dialect(JdbcChatMemoryRepositoryDialect.from(dataSource)) // detected from the JDBC URL
                .transactionManager(txManager)
                .build();
    }

    @Bean
    ChatMemory chatMemory(JdbcChatMemoryRepository repository) {
        return MessageWindowChatMemory.builder().chatMemoryRepository(repository).maxMessages(20).build();
    }
}

package com.example.docsassistant;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.cassandra.CassandraChatMemoryRepository;
import org.springframework.ai.chat.memory.repository.mongo.MongoChatMemoryRepository;
import org.springframework.ai.chat.memory.repository.neo4j.Neo4jChatMemoryRepository;
import org.springframework.ai.chat.memory.repository.redis.RedisChatMemoryRepository;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Profile;

/** Each starter auto-configures its repository bean; the window is the same code whatever the store. */
@Configuration(proxyBeanMethods = false)
class StoreBackedMemory {

    @Bean
    @Profile("redis")
    ChatMemory redisMemory(RedisChatMemoryRepository repository) {
        return window(repository);
    }

    @Bean
    @Profile("mongo")
    ChatMemory mongoMemory(MongoChatMemoryRepository repository) {
        return window(repository);
    }

    @Bean
    @Profile("neo4j")
    ChatMemory neo4jMemory(Neo4jChatMemoryRepository repository) {
        return window(repository);
    }

    @Bean
    @Profile("cassandra")
    ChatMemory cassandraMemory(CassandraChatMemoryRepository repository) {
        return window(repository);
    }

    private static ChatMemory window(org.springframework.ai.chat.memory.ChatMemoryRepository repository) {
        return MessageWindowChatMemory.builder().chatMemoryRepository(repository).maxMessages(20).build();
    }
}

package com.example.docsassistant;

import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;

/** A custom repository: conversations live in memory and expire after a period without writes. */
class InMemoryTtlChatMemoryRepository implements ChatMemoryRepository {

    private record Entry(List<Message> messages, Instant writtenAt) { }

    private final Map<String, Entry> conversations = new ConcurrentHashMap<>();
    private final Duration ttl;
    private final Clock clock;

    InMemoryTtlChatMemoryRepository(Duration ttl, Clock clock) {
        this.ttl = ttl;
        this.clock = clock;
    }

    @Override
    public List<String> findConversationIds() {
        evictExpired();
        return new ArrayList<>(this.conversations.keySet());
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        evictExpired();
        Entry entry = this.conversations.get(conversationId);
        return entry == null ? List.of() : entry.messages();
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        // the contract: this is the complete new content of the conversation, not an append
        List<Message> storable = messages.stream()
                .filter(m -> !(m instanceof ToolResponseMessage)
                        && !(m instanceof AssistantMessage a && a.hasToolCalls())) // this store keeps text only
                .toList();
        this.conversations.put(conversationId, new Entry(storable, this.clock.instant()));
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        this.conversations.remove(conversationId);
    }

    private void evictExpired() {
        Instant cutoff = this.clock.instant().minus(this.ttl);
        this.conversations.values().removeIf(entry -> entry.writtenAt().isBefore(cutoff));
    }
}

package com.example.docsassistant;

import java.sql.Timestamp;
import java.time.Duration;
import java.time.Instant;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.core.simple.JdbcClient;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

/** Deletes conversations that have been idle longer than the retention period (needs @EnableScheduling). */
@Component
class MemoryRetentionJob {

    private static final Logger log = LoggerFactory.getLogger(MemoryRetentionJob.class);

    private final JdbcClient jdbc;
    private final Duration retention = Duration.ofDays(30);

    MemoryRetentionJob(JdbcClient jdbc) {
        this.jdbc = jdbc;
    }

    @Scheduled(cron = "0 15 3 * * *")
    void purgeIdleConversations() {
        Timestamp cutoff = Timestamp.from(Instant.now().minus(this.retention));
        int rows = this.jdbc.sql("""
                        DELETE FROM spring_ai_chat_memory
                        WHERE conversation_id IN (
                            SELECT conversation_id FROM spring_ai_chat_memory
                            GROUP BY conversation_id
                            HAVING MAX("timestamp") < :cutoff)""")
                .param("cutoff", cutoff)
                .update();
        this.jdbc.sql("""
                        DELETE FROM docs_conversation
                        WHERE conversation_id NOT IN (SELECT conversation_id FROM spring_ai_chat_memory)""")
                .update();
        log.info("retention purge removed {} message rows", rows);
    }
}

CREATE TABLE IF NOT EXISTS docs_conversation (
    conversation_id VARCHAR(36) PRIMARY KEY,
    owner           VARCHAR(255) NOT NULL,
    created_at      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS docs_conversation_owner_idx ON docs_conversation(owner);

package com.example.docsassistant;

import java.util.List;

import org.springframework.jdbc.core.simple.JdbcClient;
import org.springframework.stereotype.Repository;

/**
 * Remembers which user owns which conversation id. The memory table only knows the id, and the id is a hash,
 * so without this table "delete everything this user said" cannot be answered.
 */
@Repository
class ConversationRegistry {

    private final JdbcClient jdbc;

    ConversationRegistry(JdbcClient jdbc) {
        this.jdbc = jdbc;
    }

    void register(String owner, String conversationId) {
        this.jdbc.sql("""
                        INSERT INTO docs_conversation (conversation_id, owner)
                        SELECT :id, :owner
                        WHERE NOT EXISTS (SELECT 1 FROM docs_conversation WHERE conversation_id = :id)""")
                .param("id", conversationId)
                .param("owner", owner)
                .update();
    }

    List<String> idsOwnedBy(String owner) {
        return this.jdbc.sql("SELECT conversation_id FROM docs_conversation WHERE owner = :owner")
                .param("owner", owner)
                .query(String.class)
                .list();
    }

    void forget(String owner) {
        this.jdbc.sql("DELETE FROM docs_conversation WHERE owner = :owner").param("owner", owner).update();
    }
}

package com.example.docsassistant;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
class MemoryErasureService {

    private final ChatMemoryRepository repository;
    private final ConversationRegistry registry;

    MemoryErasureService(ChatMemoryRepository repository, ConversationRegistry registry) {
        this.repository = repository;
        this.registry = registry;
    }

    /** Right of erasure: remove every conversation the user owns, and nobody else's. */
    @Transactional
    int eraseAllFor(String owner) {
        List<String> ids = this.registry.idsOwnedBy(owner);
        ids.forEach(this.repository::deleteByConversationId);
        this.registry.forget(owner);
        return ids.size();
    }

    /** Right of access: hand the user a copy of what is stored about them. */
    Map<String, List<Message>> exportFor(String owner) {
        Map<String, List<Message>> export = new LinkedHashMap<>();
        for (String id : this.registry.idsOwnedBy(owner)) {
            export.put(id, this.repository.findByConversationId(id));
        }
        return export;
    }
}

package com.example.docsassistant;

import java.util.Map;

import org.springframework.http.ResponseEntity;
import org.springframework.security.core.Authentication;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
class MemoryPrivacyController {

    private final MemoryErasureService erasure;

    MemoryPrivacyController(MemoryErasureService erasure) {
        this.erasure = erasure;
    }

    // the owner is always the caller: there is deliberately no user id in the URL or the body
    @DeleteMapping("/api/me/conversations")
    ResponseEntity<Map<String, Integer>> eraseMine(Authentication authentication) {
        return ResponseEntity.ok(Map.of("deleted", this.erasure.eraseAllFor(authentication.getName())));
    }
}