Implementing Conversation Memory in Java AI Agents for Reimbursement Systems

Core Objective: Adding Memory to Reimbursement Processes

In reimbursement scenarios, convesrations typically span multiple turns, such as:

User: "I need to expense a business trip"
Assistant: "Sure, please provide the amount and purpose"
User: "128 RMB for taxi to the airport to meet a client" (Without memory, the assistant can't connect this to "business trip")

The role of Conversation Memory: It allows the AI assistant to remember information from previous turns (in this case, the expense type is "business trip"), enabling it to correctly interpret subsequent inputs.

Key Concepts of Conversation Memory in Reimbursement Contexts

1. Memory vs. History

• History: Complete conversation records displayed to users.
• Memory: Filtered information provided to the LLM to "remember" key details.

In reimbursement applications, memory only retains essential information like expense type, amount, and purpose, omitting irrelevant conversation content.

2. Memory Implementation Approaches

Practical Implementation: Adding Memory to a Reimbursement Assistant

Step 1: Define the Memory-Enabled Assistant Interface

public interface ExpenseProcessingAssistant {
    
    /**
     * Single-user conversation method
     */
    String processRequest(String userMessage);
    
    /**
     * Multi-session isolation method (supports multiple concurrent users)
     * @param sessionId Session identifier (e.g., user ID)
     * @param userMessage User input
     */
    String processRequest(@SessionId String sessionId, @UserInput String userMessage);
}

Note: @SessionId is an annotation provided by LangChain4j to distinguish conversation memories for different users.

Step 2: Build the Memory-Enabled Assistant

Option A: Single-User Version (Quick Start)

@Configuration
public class ExpenseAIConfiguration {
    
    @Bean
    public ExpenseProcessingAssistant expenseProcessingAssistant(
            ChatLanguageModel languageModel,
            ExpenseTools expenseTools,
            ValidationTools validationTools) {
        
        // 1. Create memory: retain the most recent 10 messages
        ConversationMemory conversationMemory = MessageWindowConversationMemory.withMaxMessages(10);
        
        // 2. Build the assistant
        return AiServices.builder(ExpenseProcessingAssistant.class)
                .chatLanguageModel(languageModel)
                .chatMemory(conversationMemory)  // Inject memory
                .tools(expenseTools, validationTools)
                .build();
    }
}

Option B: Multi-User Version (Production Recommended)

@Configuration
public class ExpenseAIConfiguration {
    
    @Bean
    public ExpenseProcessingAssistant expenseProcessingAssistant(
            ChatLanguageModel languageModel,
            ExpenseTools expenseTools,
            ValidationTools validationTools) {
        
        // Use ConversationMemoryProvider: create independent memory for each sessionId
        ConversationMemoryProvider memoryProvider = sessionId -> 
            MessageWindowConversationMemory.builder()
                .id(sessionId)  // Use sessionId to differentiate users
                .maxMessages(10)
                .build();
        
        return AiServices.builder(ExpenseProcessingAssistant.class)
                .chatLanguageModel(languageModel)
                .chatMemoryProvider(memoryProvider)  // Inject memory provider
                .tools(expenseTools, validationTools)
                .build();
    }
}

Important: The ConversationMemoryProvider automatically creates independent memory instances for each sessionId, preventing interference between users.

Step 3: Verify Memory Functionality

@SpringBootTest
class ExpenseProcessingAssistantTest {
    
    @Autowired
    private ExpenseProcessingAssistant assistant;
    
    @Test
    void testMultiRoundExpenseProcessing() {
        // First round: User initiates expense request
        String response1 = assistant.processRequest("user-123", "I want to claim taxi expenses");
        System.out.println(response1);
        // Expected output: "Please provide the taxi amount and purpose"
        
        // Second round: User provides additional information (assistant should remember "taxi expenses")
        String response2 = assistant.processRequest("user-123", "128 RMB for taxi to the airport to meet a client");
        System.out.println(response2);
        // Expected output: "Recorded your taxi expense: 128 RMB for traveling to the airport to meet a client. Is there anything else I can help with?"
        
        // Test isolation for another user
        String responseOther = assistant.processRequest("user-456", "128 RMB for taxi to the airport to meet a client");
        // Expected output: "Please tell me what type of expense you're claiming, as there's no context"
    }
}

Upgrading to TokenWindowConversationMemory

The core difference between TokenWindowConversationMemory and MessageWindowConversationMemory lies in their message eviction strategies:

• MessageWindowConversationMemory: Retains the most recent N messages based on count.
• TokenWindowConversationMemory: Retains messages based on token count. When the total exceeds the limit, the oldest messages are automatically discarded until the constraint is satisfied.

Why Upgrade?

• Precise context size control: Different messages vary significantly in length (e.g., a single sentence vs. a code block). Retaining 10 messages by count might result in contexts of just a few tokens for some conversations, while exceeding model limits (e.g., 4096) for others. TokenWindowConversationMemory ensures the total token count never exceeds your specified limit, preventing truncation or overflow errors.
• Token cost savings (API expenses): For token-based billing models like OpenAI's, shorter contexts mean lower costs per API call. TokenWindowConversationMemory allows you to proactively control the number of tokens sent to the model, avoiding unnecessary expenses.
• Potential memory space savings: If messages are typically long, retaining a fixed token count will result in fewer message objects being stored, reducing memory usage. Conversely, if messages are short, more messages might be retained compared to a count-based strategy, potentially increasing memory usage. However, token consumption optimization is a clear benefit.

How to Migrate from MessageWindowConversationMemory to TokenWindowConversationMemory

1. Add Tokenizer Dependency

TokenWindowConversationMemory needs to know how to count text tokens, requiring a Tokenizer implementation. For OpenAI:

<dependency>
    <groupid>dev.langchain4j</groupid>
    <artifactid>langchain4j-open-ai</artifactid>
    <version>1.11.0</version>
</dependency>

2. Modify Configuration to Use TokenWindowConversationMemory

Original configuration using MessageWindowConversationMemory:

ConversationMemoryProvider memoryProvider = sessionId ->
    MessageWindowConversationMemory.builder()
        .id(sessionId)
        .maxMessages(10)
        .build();

Changed to use TokenWindowConversationMemory:

        // Create Tokenizer (recommended as singleton)
        OpenAiTokenEstimator tokenEstimator = new OpenAiTokenEstimator(modelName);
        // Use tokenMemoryProvider: create independent memory for each sessionId
        ConversationMemoryProvider tokenMemoryProvider = sessionId ->
                TokenWindowConversationMemory.builder()
                        .id(sessionId)  // Use sessionId to differentiate users
                        .maxTokens(5000, tokenEstimator)
                        .build();

Recommendations for setting maxTokens:

Consider the model's maximum context limit (e.g., 4096 for GPT-3.5). Reserve space for the model's response (e.g., 1000 tokens), so the historical token count should be set to modelMaxTokens - maxCompletionTokens - safety margin.

For example: Model limit 4096, reserve 1000 for response, safety margin 200, then maxTokens = 4096 - 1000 - 200 = 2896.

3. Complete Example (with Persistence)

If you previously implemented ConversationMemoryStorage (e.g., Redis), simply replace the ConversationMemory construction logic while keeping ConversationMemoryStorage unchanged (it still stores all messages, with the window strategy determining which messages are provided to the LLM).

    @Bean
    public UnifiedExpenseAssistant unifiedExpenseAssistant(StreamingChatModel streamingModel,
                                             ExpenseTools expenseTools,
                                             ValidationTools validationTools,
                                             ExpenseFlowTool expenseFlowTool) {
        // MessageWindowConversationMemory: retain the most recent N messages (sliding window)
        // ConversationMemory memory = MessageWindowConversationMemory.withMaxMessages(10);
        
        // TokenWindowConversationMemory: retain the most recent N tokens (sliding window)
        OpenAiTokenEstimator tokenEstimator = new OpenAiTokenEstimator(modelName);
        // Use ConversationMemoryProvider: create independent memory for each sessionId
        ConversationMemoryProvider memoryProvider = sessionId ->
                MessageWindowConversationMemory.builder()
                        .id(sessionId)  // Use sessionId to differentiate users
                        .maxMessages(10)
                        .build();
        // Use tokenMemoryProvider: create independent memory for each sessionId
        ConversationMemoryProvider tokenMemoryProvider = sessionId ->
                TokenWindowConversationMemory.builder()
                        .id(sessionId)  // Use sessionId to differentiate users
                        .maxTokens(5000, tokenEstimator)
                        .build();
        // Create AI service instance using AiServices
        return AiServices.builder(UnifiedExpenseAssistant.class)
                .streamingChatModel(streamingModel)
                // .chatMemory(memory)
                // .chatMemoryProvider(memoryProvider) // Store by conversation
                .chatMemoryProvider(tokenMemoryProvider)
                // Key point: pass multiple tool instances through .tools() method
                .tools(expenseTools, validationTools, expenseFlowTool)
                .build();
    }

Does TokenWindowConversationMemory Save Space Compared to MessageWindowConversationMemory?

Memory Space (JVM Memory)

• MessageWindowConversationMemory: Stores a fixed number of ChatMessage objects, each containing role, content, etc. Memory usage is proportional to messsage length.
• TokenWindowConversationMemory: Stores a variable number of messages depending on token length. When messages are long, fewer messages are retained compared to a fixed count, saving memory. When messages are short, more messages might be retained, potentially increasing memory usage.

Therefore, we cannot definitively say TokenWindow always saves memory, but it allows you to control the memory ceiling based on token budget (since the total token count is limited, the number of message objects also has an upper bound).

Persistent Storage (e.g., Redis)

Regardless of the window strategy, if you use ConversationMemoryStorage, all historical messages are stored (unless actively deleted). The window strategy only determines which messages are loaded from persistent storage and provided to the LLM, without affecting the total size of persistent storage. If you want persistent storage to also be pruned by token count, you would need to implement pruning logic in updateMessages, but this is generally not recommended as the user interface may require complete history.

Token Cost

This is the core advantage of TokenWindowConversationMemory: It precisely controls token consumption per request, avoiding overages or waste due to excessive messages. In long-running systems, it can significantly reduce API costs.

Testing and Validation

After upgrading, you can verify the functionality through these methods:

• Check actual token count: Enable LangChain4j logging (logging.level.dev.langchain4j=DEBUG) and observe token statistics for each request.
• Simulate long conversations: Send multiple long messages and confirm that when total tokens approach maxTokens, the earliest messages are discarded.
• Check memory usage (optional): Use JVM memory analysis tools to compare memory changes between the two strategies.

Comparison Summary

Recommendation: Prioritize TokenWindowConversationMemory in production environments, combined with persistent storage, to achieve the best balance between cost and performance.