Spring AI Chat Memory

Spring AI - Chat Memory & Conversation Management

Description

Comprehensive guide for chat memory implementation in Spring AI. Covers memory storage strategies, conversation management, token-aware windowing, supported backends (in-memory, JDBC, Cassandra, MongoDB, etc.), and integration with chat advisors.

When to Use

• Maintaining conversation history
• Multi-turn chat applications
• Stateful AI conversations
• User session management
• Token budget management
• Persistent conversation storage
• Conversation analytics
• Multi-user chat systems

Topics Covered

1. Chat Memory Concepts

• Message History: Preserving past messages
• Context Window: Fitting messages in token limits
• Windowing Strategies: What to keep/discard
• Metadata: Timestamps, user info, turn markers
• Cleanup: Managing memory size over time

2. Memory Implementations

InMemoryChatMemoryStore

• Default, fast, no external dependency
• Lost on application restart
• Suitable for testing/development
• Thread-safe operations
• Single-node only

JdbcChatMemoryStore

• Persistent SQL database storage
• Works with PostgreSQL, MySQL, Oracle, etc.
• JDBC abstraction for portability
• Transaction support
• Scalable to many users

CassandraChatMemoryStore

• Distributed NoSQL storage
• High availability and scalability
• Time-series optimized
• Partition key: conversation ID
• Row key: message sequence

MongoDbChatMemoryStore

• Document-based storage
• Natural JSON representation
• Flexible schema
• Built-in replication
• Query flexibility

Neo4jChatMemoryStore

• Graph-based storage
• Relationship tracking
• Conversation entity relationships
• Pattern-based queries
• Knowledge graph integration

AzureCosmosDbChatMemoryStore

• Azure serverless option
• Multi-region replication
• Partition by conversation
• Serverless scalability
• Global distribution

3. Message Management

• Adding messages: User + Assistant exchanges
• Listing messages: Range queries, filtering
• Deleting messages: Cleanup, privacy
• Updating metadata: Annotations, corrections
• Purging: Conversation cleanup

4. Window Management

• MessageWindowChatMemoryStore: Token-aware windowing
• Fixed window: Last N messages
• Token-based window: Fit within token budget
• Time-based window: Recent N minutes
• Hybrid window: Combine strategies

5. Memory Advisor Integration

• Automatic memory attachment to requests
• Message injection into chat history
• Conversation ID tracking
• User ID tracking
• Memory callback patterns

Code Patterns

In-Memory Chat Memory (Default)

@Configuration
public class ChatMemoryConfig {
    @Bean
    public ChatMemoryStore chatMemoryStore() {
        return new InMemoryChatMemoryStore();
    }
}

JDBC-based Persistent Memory

@Configuration
public class PersistentMemoryConfig {
    @Bean
    public ChatMemoryStore chatMemoryStore(JdbcTemplate jdbc) {
        return new JdbcChatMemoryStore(jdbc);
    }
}

MongoDB Chat Memory

@Configuration
public class MongoDbMemoryConfig {
    @Bean
    public ChatMemoryStore chatMemoryStore(
            MongoTemplate mongoTemplate) {
        return new MongoDbChatMemoryStore(mongoTemplate);
    }
}

Chat Memory with Token Windowing

@Configuration
public class SmartMemoryConfig {
    @Bean
    public ChatMemoryStore chatMemoryStore() {
        InMemoryChatMemoryStore store = new InMemoryChatMemoryStore();

        // Token-aware windowing (max 2000 tokens)
        return new MessageWindowChatMemoryStore(
            store,
            2000,  // max tokens
            chatModel,  // for token counting
            EmbeddingModel.DEFAULT_TOKEN_COUNTER
        );
    }
}

Using Chat Memory with ChatClient

@Service
public class ConversationService {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatMemoryStore memoryStore;

    public String chat(String userId, String conversationId, String message) {
        // Store user message
        memoryStore.add(conversationId,
            new Message(MessageType.USER, message));

        // Get chat history
        List<Message> history = memoryStore.get(conversationId);

        // Call LLM with context
        String response = chatClient.prompt()
            .messages(history)  // Include history
            .user(message)
            .call()
            .content();

        // Store assistant response
        memoryStore.add(conversationId,
            new Message(MessageType.ASSISTANT, response));

        return response;
    }
}

Chat Memory Advisor Pattern

@Configuration
public class MemoryAdvisorConfig {
    @Bean
    public ChatClientRequestAdvisor chatMemoryAdvisor(
            ChatMemoryStore memoryStore) {

        return ChatMemoryAdvisor.builder()
            .chatMemory(memoryStore)
            .userIdResolver(request -> extractUserId(request))
            .conversationIdResolver(request -> extractConversationId(request))
            .build();
    }
}

@Service
public class SmartChat {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatClientRequestAdvisor memoryAdvisor;

    public String chat(String message, String conversationId) {
        // Memory advisor automatically:
        // 1. Loads previous messages
        // 2. Injects into context
        // 3. Stores new message

        return chatClient.prompt()
            .user(message)
            .advisors(memoryAdvisor)
            .call()
            .content();
    }
}

Message Window Management

@Service
public class WindowedMemoryService {
    @Autowired
    private ChatMemoryStore memoryStore;

    public List<Message> getMemoryWindow(
            String conversationId,
            int maxTokens) {

        List<Message> allMessages = memoryStore.get(conversationId);

        // Calculate tokens and create window
        int tokenCount = 0;
        List<Message> window = new ArrayList<>();

        // Include from most recent backwards
        for (int i = allMessages.size() - 1; i >= 0; i--) {
            Message msg = allMessages.get(i);
            int msgTokens = countTokens(msg.getContent());

            if (tokenCount + msgTokens <= maxTokens) {
                window.add(0, msg);  // prepend
                tokenCount += msgTokens;
            } else {
                break;
            }
        }

        return window;
    }
}

Conversation Management

@Service
public class ConversationManager {
    @Autowired
    private ChatMemoryStore memoryStore;

    public void startConversation(String conversationId, String initialContext) {
        // Store system context
        memoryStore.add(conversationId,
            new Message(MessageType.SYSTEM, initialContext));
    }

    public void endConversation(String conversationId) {
        // Preserve for analytics, then cleanup
        archiveConversation(conversationId);
        memoryStore.delete(conversationId);
    }

    public List<Message> getConversation(String conversationId) {
        return memoryStore.get(conversationId);
    }

    public void clearMemory(String conversationId) {
        memoryStore.delete(conversationId);
    }
}

Multi-turn Conversation Example

@Service
public class MultiTurnChat {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatMemoryAdvisor memoryAdvisor;

    public String runConversation(String conversationId) {
        List<String> interactions = List.of(
            "What is Spring Framework?",
            "Tell me more about Spring Boot",
            "How does Spring AI relate to these?"
        );

        for (String userInput : interactions) {
            String response = chatClient.prompt()
                .user(userInput)
                .advisors(memoryAdvisor)
                .call()
                .content();

            logger.info("User: {}", userInput);
            logger.info("Assistant: {}", response);
        }

        // Memory advisor handled history automatically
        return "Conversation complete";
    }
}

Configuration

JDBC Storage Setup

@Configuration
public class JdbcMemoryConfig {
    @Bean
    public ChatMemoryStore chatMemoryStore(JdbcTemplate jdbc) {
        JdbcChatMemoryStore store = new JdbcChatMemoryStore(jdbc);
        store.createSchema();  // Auto-create tables
        return store;
    }
}

Properties

# In-Memory (default)
spring.ai.memory.store=in-memory

# JDBC
spring.ai.memory.store=jdbc
spring.ai.memory.jdbc.table-name=chat_memory
spring.ai.memory.jdbc.create-table=true

# Cassandra
spring.ai.memory.store=cassandra
spring.ai.memory.cassandra.keyspace=ai_chat
spring.ai.memory.cassandra.table-name=messages

# MongoDB
spring.ai.memory.store=mongodb
spring.ai.memory.mongodb.collection=conversations

# Window settings
spring.ai.memory.window.strategy=token-based
spring.ai.memory.window.max-tokens=2000
spring.ai.memory.window.max-messages=50
spring.ai.memory.window.time-to-live=3600

Storage Schema Examples

JDBC Storage

CREATE TABLE chat_memory (
    id VARCHAR(36) PRIMARY KEY,
    conversation_id VARCHAR(255) NOT NULL,
    user_id VARCHAR(255),
    message_type VARCHAR(20),
    content TEXT,
    metadata JSONB,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    INDEX idx_conversation_id (conversation_id),
    INDEX idx_user_id (user_id)
);

MongoDB Storage

{
    "_id": ObjectId,
    "conversation_id": "conv-123",
    "user_id": "user-456",
    "messages": [
        {
            "role": "user",
            "content": "...",
            "timestamp": ISODate()
        }
    ]
}

Best Practices

• Always set appropriate memory limits
• Implement conversation cleanup policies
• Use token-based windowing for efficiency
• Backup important conversations
• Monitor memory storage growth
• Implement user privacy controls
• Archive old conversations
• Test memory persistence

Related Skills

• advisors/SKILL.md - Memory advisor integration
• chat-models/SKILL.md - Chat operations
• observability/SKILL.md - Monitoring storage

References

• API: /pages/api/chat-memory.adoc
• Advisors: /pages/api/advisors.adoc
• Database integration: Provider-specific docs