improve app, refactor, add comment

Author: Yuriy Bezsonov

samples/spring-ai-te-agent/ai-agent/src/main/java/com/example/ai/agent/controller/ChatController.java

@@ -5,10 +5,7 @@
 import com.example.ai.agent.service.ConversationSummaryService;
 import com.example.ai.agent.service.DocumentChatService;
 import org.springframework.http.MediaType;
-import org.springframework.web.bind.annotation.PostMapping;
-import org.springframework.web.bind.annotation.RequestBody;
-import org.springframework.web.bind.annotation.RequestMapping;
-import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.bind.annotation.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
@@ -39,11 +36,10 @@ public Flux<String> chat(@RequestBody ChatRequest request, Principal principal)
         String userId = getUserId(request.userId(), principal);
         chatMemoryService.setCurrentUserId(userId);
 
-        if (hasFile(request)) {
-            return documentChatService.processDocument(request.prompt(), request.fileBase64(), request.fileName());
-        } else {
-            return chatService.processChat(request.prompt());
-        }
+        // Route to document analysis or regular chat
+        return hasFile(request)
+            ? documentChatService.processDocument(request.prompt(), request.fileBase64(), request.fileName())
+            : chatService.processChat(request.prompt());
     }
 
     @PostMapping("summarize")
@@ -58,9 +54,11 @@ public String summarize(@RequestBody(required = false) SummarizeRequest request,
     }
 
     private String getUserId(String requestUserId, Principal principal) {
+        // Production: use authenticated user from Spring Security
         if (principal != null) {
             return principal.getName();
         }
+        // Development: use provided userId or default
         return requestUserId != null ? requestUserId : "user1";
     }
 

samples/spring-ai-te-agent/ai-agent/src/main/java/com/example/ai/agent/controller/WebViewController.java

@@ -16,6 +16,7 @@ public class WebViewController {
 
     @GetMapping("/")
     public String index(Model model) {
+        // Pass feature flags to UI for conditional rendering
         model.addAttribute("multiUserEnabled", multiUserEnabled);
         model.addAttribute("multiModalEnabled", multiModalEnabled);
         return "chat";

samples/spring-ai-te-agent/ai-agent/src/main/java/com/example/ai/agent/service/ChatMemoryService.java

@@ -4,15 +4,10 @@
 import org.springframework.ai.chat.memory.MessageWindowChatMemory;
 import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
 import org.springframework.ai.chat.memory.repository.jdbc.PostgresChatMemoryRepositoryDialect;
-import org.springframework.ai.chat.messages.UserMessage;
-import org.springframework.ai.chat.messages.AssistantMessage;
-import org.springframework.ai.chat.messages.SystemMessage;
-import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.*;
 import org.springframework.ai.chat.prompt.Prompt;
 import reactor.core.publisher.Flux;
-
 import org.springframework.stereotype.Service;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import javax.sql.DataSource;
@@ -25,15 +20,14 @@ public class ChatMemoryService {
     private static final int MAX_CONTEXT_SUMMARIES = 10;
     private static final int MAX_PREFERENCES = 1;
 
+    // Three-tier memory: Session (recent), Context (summaries), Preferences (profile)
     private final MessageWindowChatMemory sessionMemory;
     private final MessageWindowChatMemory contextMemory;
     private final MessageWindowChatMemory preferencesMemory;
-
-    // Thread-local to store current userId per request
     private final ThreadLocal<String> currentUserId = ThreadLocal.withInitial(() -> "user1");
 
     public ChatMemoryService(DataSource dataSource) {
-
+        // Single JDBC repository shared by all three memory tiers
         var jdbcRepository = JdbcChatMemoryRepository.builder()
             .dataSource(dataSource)
             .dialect(new PostgresChatMemoryRepositoryDialect())
@@ -55,79 +49,50 @@ public ChatMemoryService(DataSource dataSource) {
             .build();
     }
 
-    public MessageWindowChatMemory getSessionMemory() {
-        return sessionMemory;
-    }
-
-    public MessageWindowChatMemory getContextMemory() {
-        return contextMemory;
-    }
-
-    public MessageWindowChatMemory getPreferencesMemory() {
-        return preferencesMemory;
-    }
-
-    public void setCurrentUserId(String userId) {
-        this.currentUserId.set(userId);
-    }
-
-    public String getCurrentConversationId() {
-        return currentUserId.get();
-    }
-
     public Flux<String> callWithMemory(ChatClient chatClient, String prompt) {
         String conversationId = getCurrentConversationId();
 
-        // Check if first message - load previous context from JDBC
+        // 1. Auto-load previous context on first message
         if (sessionMemory.get(conversationId).isEmpty()) {
             loadPreviousContext(conversationId, chatClient);
         }
 
-        // Add user message to session memory
-        UserMessage userMessage = new UserMessage(prompt);
-        sessionMemory.add(conversationId, userMessage);
+        // 2. Add user message to session memory
+        sessionMemory.add(conversationId, new UserMessage(prompt));
 
-        // Stream response with conversation history
+        // 3. Stream AI response with full conversation history
         StringBuilder fullResponse = new StringBuilder();
-
         return chatClient
                 .prompt(new Prompt(sessionMemory.get(conversationId)))
                 .stream()
                 .content()
-                .doOnNext(chunk -> {
-                    fullResponse.append(chunk);
-                    logger.debug("Streaming chunk: {} chars", chunk.length());
-                })
+                .doOnNext(fullResponse::append)
                 .doOnComplete(() -> {
-                    // Add complete assistant response to session memory after streaming completes
+                    // 4. Save complete response to session memory
                     String responseText = fullResponse.toString();
-                    if (responseText != null && !responseText.isEmpty()) {
-                        AssistantMessage assistantMessage = new AssistantMessage(responseText);
-                        sessionMemory.add(conversationId, assistantMessage);
-                        logger.info("Added assistant response to memory: {} chars", responseText.length());
+                    if (!responseText.isEmpty()) {
+                        sessionMemory.add(conversationId, new AssistantMessage(responseText));
+                        logger.info("Saved response to memory: {} chars", responseText.length());
                     }
-                    logger.info("Completed streaming response.");
                 });
     }
 
     private void loadPreviousContext(String conversationId, ChatClient chatClient) {
-        logger.info("Loading previous context for conversation: {}", conversationId);
+        logger.info("Loading previous context for: {}", conversationId);
 
-        // Load user preferences (userId_preferences)
+        // 1. Load preferences (userId_preferences) and context (userId_context)
         List<Message> preferences = preferencesMemory.get(conversationId + "_preferences");
         String preferencesText = preferences.isEmpty() ? "" : preferences.get(0).getText();
-
-        // Load context summaries (userId_context)
         List<Message> summaries = contextMemory.get(conversationId + "_context");
 
         if (summaries.isEmpty() && preferencesText.isEmpty()) {
             logger.info("No previous context found");
             return;
         }
 
-        logger.info("Found {} context summaries and {} preferences", summaries.size(), preferences.isEmpty() ? 0 : 1);
+        logger.info("Found {} summaries, {} preferences", summaries.size(), preferences.isEmpty() ? 0 : 1);
 
-        // Combine preferences and summaries
+        // 2. Combine preferences and summaries
         StringBuilder contextBuilder = new StringBuilder();
         if (!preferencesText.isEmpty()) {
             contextBuilder.append("User Preferences:\n").append(preferencesText).append("\n\n");
@@ -137,13 +102,15 @@ private void loadPreviousContext(String conversationId, ChatClient chatClient) {
             summaries.forEach(msg -> contextBuilder.append(msg.getText()).append("\n\n"));
         }
 
-        String combinedContext = contextBuilder.toString();
+        // 3. Summarize combined context with AI
         var chatResponse = chatClient.prompt()
-            .user("Summarize this user information concisely:\n\n" + combinedContext)
+            .user("Summarize this user information concisely:\n\n" + contextBuilder)
             .call()
             .chatResponse();
 
-        String contextSummary = (chatResponse != null && chatResponse.getResult() != null && chatResponse.getResult().getOutput() != null)
+        String contextSummary = (chatResponse != null &&
+            chatResponse.getResult() != null &&
+            chatResponse.getResult().getOutput() != null)
             ? chatResponse.getResult().getOutput().getText()
             : null;
 
@@ -152,13 +119,34 @@ private void loadPreviousContext(String conversationId, ChatClient chatClient) {
             return;
         }
 
-        logger.info("Loaded context summary: {}", contextSummary);
-
+        // 4. Add context as system message to session memory
         String contextMessage = String.format(
             "You are continuing a conversation with this user. Here is what you know:\n\n%s\n\n" +
             "Use this information to provide personalized responses.",
             contextSummary
         );
         sessionMemory.add(conversationId, new SystemMessage(contextMessage));
+        logger.info("Loaded context summary");
+    }
+
+    // Getters and setters
+    public MessageWindowChatMemory getSessionMemory() {
+        return sessionMemory;
+    }
+
+    public MessageWindowChatMemory getContextMemory() {
+        return contextMemory;
+    }
+
+    public MessageWindowChatMemory getPreferencesMemory() {
+        return preferencesMemory;
+    }
+
+    public void setCurrentUserId(String userId) {
+        this.currentUserId.set(userId);
+    }
+
+    public String getCurrentConversationId() {
+        return currentUserId.get();
     }
 }

samples/spring-ai-te-agent/ai-agent/src/main/java/com/example/ai/agent/service/ChatService.java

@@ -4,7 +4,6 @@
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Service;
 import reactor.core.publisher.Flux;
-
 import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
 import org.springframework.ai.vectorstore.VectorStore;
@@ -36,20 +35,19 @@ public ChatService(ChatMemoryService chatMemoryService,
                       ToolCallbackProvider tools,
                       ChatClient.Builder chatClientBuilder) {
 
+        // Build ChatClient with RAG, tools, and memory
         this.chatClient = chatClientBuilder
                 .defaultSystem(SYSTEM_PROMPT)
-                .defaultAdvisors(
-                        QuestionAnswerAdvisor.builder(vectorStore).build()
-                )
-                .defaultTools(dateTimeService, weatherService)
-                .defaultToolCallbacks(tools)
+                .defaultAdvisors(QuestionAnswerAdvisor.builder(vectorStore).build())  // RAG for policies
+                .defaultTools(dateTimeService, weatherService)  // Custom tools
+                .defaultToolCallbacks(tools)  // Auto-registered tools from MCP
                 .build();
 
         this.chatMemoryService = chatMemoryService;
     }
 
     public Flux<String> processChat(String prompt) {
-        logger.info("Processing streaming chat request - prompt: '{}'", prompt);
+        logger.info("Processing chat: '{}'", prompt);
         try {
             // Simple streaming without memory:
             // return chatClient
@@ -58,7 +56,7 @@ public Flux<String> processChat(String prompt) {
             //     .content();
             return chatMemoryService.callWithMemory(chatClient, prompt);
         } catch (Exception e) {
-            logger.error("Error processing streaming chat request", e);
+            logger.error("Error processing chat", e);
             return Flux.just("I don't know - there was an error processing your request.");
         }
     }