sonpx/loyalty-agent-service
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: overhaul chat architecture with persistent JdbcChatMemory, reactive UI components, and improved conversation management

Browse Source
This commit is contained in:
Son Phung
2026-07-23 19:20:28 +07:00
parent 2bbbff1ec9
commit 506008c69f
79 changed files with 5058 additions and 1193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/LoyaltyAgentApplication.java
View File

@@ -10,3 +10,5 @@ public class LoyaltyAgentApplication {
SpringApplication.run(LoyaltyAgentApplication.class, args);
}
}

25
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/config/FlywayConfig.java Normal file
View File

@@ -0,0 +1,25 @@
package dev.sonpx.loyalty.agent.config;
import org.flywaydb.core.Flyway;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.context.annotation.Configuration;
import org.springframework.lang.NonNull;
import javax.sql.DataSource;
@Configuration
public class FlywayConfig implements BeanPostProcessor {
@Override
public Object postProcessAfterInitialization(@NonNull Object bean, @NonNull String beanName) throws BeansException {
if (bean instanceof DataSource) {
System.out.println("--- Running Flyway Migration before DataSource is used ---");
Flyway.configure()
.dataSource((DataSource) bean)
.locations("classpath:db/migration")
.load()
.migrate();
}
return bean;
}
}

96
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/controller/ConversationController.java
View File

@@ -2,15 +2,21 @@ package dev.sonpx.loyalty.agent.controller;
import dev.sonpx.loyalty.agent.domain.Conversation;
import dev.sonpx.loyalty.agent.domain.Message;
import dev.sonpx.loyalty.agent.dto.ConversationSummaryDto;
import dev.sonpx.loyalty.agent.repository.ConversationRepository;
import dev.sonpx.loyalty.agent.repository.MessageRepository;
import org.springframework.http.ResponseEntity;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.bind.annotation.*;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
@RestController
@RequestMapping("/api/v1/conversations")
@CrossOrigin(origins = "*")
@@ -24,10 +30,56 @@ public class ConversationController {
this.messageRepository = messageRepository;
}
@GetMapping
@Transactional
public ResponseEntity<List<ConversationSummaryDto>> getConversations() {
List<Conversation> conversations = conversationRepository.findAllByOrderByUpdatedAtDesc();
List<ConversationSummaryDto> dtos = new ArrayList<>();
for (Conversation c : conversations) {
List<Message> messages = messageRepository.findByConversationIdOrderByCreatedAtAsc(c.getId());
// Clean up empty conversations from DB (except if it's the only one and newly created)
if (messages.isEmpty()) {
conversationRepository.delete(c);
continue;
}
String title = c.getTitle();
if (title == null || title.isBlank() || "Cuộc trò chuyện mới".equalsIgnoreCase(title)) {
Optional<Message> firstUserMsg = messages.stream()
.filter(m -> "USER".equalsIgnoreCase(m.getRole()))
.findFirst();
if (firstUserMsg.isPresent()) {
String content = firstUserMsg.get().getContent();
title = content.length() > 40 ? content.substring(0, 40) + "..." : content;
} else {
title = "Cuộc trò chuyện mới";
}
}
String preview = "";
Message lastMsg = messages.getLast();
String content = lastMsg.getContent();
preview = content.length() > 100 ? content.substring(0, 100) + "..." : content;
dtos.add(ConversationSummaryDto.builder()
.id(c.getId())
.title(title)
.preview(preview)
.createdAt(c.getCreatedAt())
.updatedAt(c.getUpdatedAt())
.build());
}
return ResponseEntity.ok(dtos);
}
@PostMapping
public ResponseEntity<Conversation> createConversation() {
Conversation conversation = new Conversation();
conversation.setId(UUID.randomUUID().toString());
conversation.setTitle("Cuộc trò chuyện mới");
conversation.setCreatedAt(Instant.now());
conversation.setUpdatedAt(Instant.now());
Conversation saved = conversationRepository.save(conversation);
@@ -39,4 +91,48 @@ public class ConversationController {
List<Message> messages = messageRepository.findByConversationIdOrderByCreatedAtAsc(id);
return ResponseEntity.ok(messages);
}
@PutMapping("/{id}")
public ResponseEntity<ConversationSummaryDto> updateTitle(
@PathVariable("id") String id,
@RequestBody Map<String, String> payload) {
String newTitle = payload.get("title");
if (newTitle == null || newTitle.isBlank()) {
return ResponseEntity.badRequest().build();
}
return conversationRepository.findById(id).map(conv -> {
conv.setTitle(newTitle.trim());
conv.setUpdatedAt(Instant.now());
Conversation saved = conversationRepository.save(conv);
List<Message> messages = messageRepository.findByConversationIdOrderByCreatedAtAsc(id);
String preview = "";
if (!messages.isEmpty()) {
String content = messages.getLast().getContent();
preview = content.length() > 100 ? content.substring(0, 100) + "..." : content;
}
return ResponseEntity.ok(ConversationSummaryDto.builder()
.id(saved.getId())
.title(saved.getTitle())
.preview(preview)
.createdAt(saved.getCreatedAt())
.updatedAt(saved.getUpdatedAt())
.build());
}).orElseGet(() -> ResponseEntity.notFound().build());
}
@DeleteMapping("/{id}")
@Transactional
public ResponseEntity<Void> deleteConversation(@PathVariable("id") String id) {
List<Message> messages = messageRepository.findByConversationIdOrderByCreatedAtAsc(id);
if (!messages.isEmpty()) {
messageRepository.deleteAll(messages);
}
conversationRepository.deleteById(id);
return ResponseEntity.noContent().build();
}
}

16
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/executor/SimpleAgentExecutor.java
View File

@@ -5,6 +5,7 @@ import dev.sonpx.loyalty.agent.core.memory.AgentMemoryManager;
import dev.sonpx.loyalty.agent.core.orchestrator.PromptBuilder;
import dev.sonpx.loyalty.agent.domain.ChatRequest;
import dev.sonpx.loyalty.agent.service.AgentEventPublisher;
import dev.sonpx.loyalty.agent.service.AgentPersistenceService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.tool.ToolCallback;
@@ -34,19 +35,22 @@ public class SimpleAgentExecutor {
private final AgentEventPublisher eventPublisher;
private final AgentMemoryManager memoryManager;
private final PromptBuilder promptBuilder;
private final AgentPersistenceService persistenceService;
public SimpleAgentExecutor(@Qualifier("queryAgentChatClient") ChatClient queryChatClient,
@Qualifier("chatAgentChatClient") ChatClient chatChatClient,
ToolInterceptor toolInterceptor,
AgentEventPublisher eventPublisher,
AgentMemoryManager memoryManager,
PromptBuilder promptBuilder) {
PromptBuilder promptBuilder,
AgentPersistenceService persistenceService) {
this.queryChatClient = queryChatClient;
this.chatChatClient = chatChatClient;
this.toolInterceptor = toolInterceptor;
this.eventPublisher = eventPublisher;
this.memoryManager = memoryManager;
this.promptBuilder = promptBuilder;
this.persistenceService = persistenceService;
}
/**
@@ -63,6 +67,9 @@ public class SimpleAgentExecutor {
log.info("SimpleAgent executing with tools for conversation: {}, connection: {}, intent: {}", conversationKey, connectionSessionId, intent);
log.debug("System prompt: {}", systemPrompt);
persistenceService.saveUserMessage(conversationKey, request.messageId(), request.prompt());
StringBuilder responseAccumulator = new StringBuilder();
queryChatClient.prompt()
.advisors(memoryManager.getAdvisor())
.advisors(a -> a
@@ -75,11 +82,13 @@ public class SimpleAgentExecutor {
.content()
.doOnComplete(() -> {
log.info("SimpleAgent execution completed for conversation: {}, connection: {}", conversationKey, connectionSessionId);
persistenceService.saveAssistantMessage(conversationKey, responseAccumulator.toString());
eventPublisher.done(connectionSessionId, request.messageId());
})
.subscribe(
content -> {
if (!content.isEmpty()) {
responseAccumulator.append(content);
eventPublisher.token(connectionSessionId, request.messageId(), content);
}
},
@@ -101,6 +110,9 @@ public class SimpleAgentExecutor {
log.info("DirectChat executing for conversation: {}, connection: {}", conversationKey, connectionSessionId);
persistenceService.saveUserMessage(conversationKey, request.messageId(), request.prompt());
StringBuilder responseAccumulator = new StringBuilder();
chatChatClient.prompt()
.advisors(memoryManager.getAdvisor())
.advisors(a -> a
@@ -112,11 +124,13 @@ public class SimpleAgentExecutor {
.content()
.doOnComplete(() -> {
log.info("DirectChat completed for conversation: {}, connection: {}", conversationKey, connectionSessionId);
persistenceService.saveAssistantMessage(conversationKey, responseAccumulator.toString());
eventPublisher.done(connectionSessionId, request.messageId());
})
.subscribe(
content -> {
if (!content.isEmpty()) {
responseAccumulator.append(content);
eventPublisher.token(connectionSessionId, request.messageId(), content);
}
},

106
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/executor/ToolInterceptor.java
View File

@@ -1,24 +1,32 @@
package dev.sonpx.loyalty.agent.core.executor;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import dev.sonpx.loyalty.agent.exception.UnrecoverableToolExecutionException;
import dev.sonpx.loyalty.agent.service.AgentEventPublisher;
import dev.sonpx.loyalty.agent.service.ToolResultPresentationProcessor;
import io.modelcontextprotocol.client.McpSyncClient;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.stereotype.Component;
import java.io.EOFException;
import java.net.ConnectException;
import java.net.SocketException;
import java.nio.channels.ClosedChannelException;
import java.util.ArrayList;
import java.util.List;
/**
* Wraps MCP tool callbacks with:
* 1. Event publishing (tool running/done status)
* 2. Error sanitization (hides stack traces from LLM)
* 3. Result truncation (prevents context window overflow)
* 4. Result presentation processing (agent_instruction handling)
* 2. Error sanitization & Circuit breaker for unrecoverable errors
* 3. Auto-reconnection on dropped SSE/MCP connection
* 4. Result truncation (prevents context window overflow)
* 5. Result presentation processing (agent_instruction handling)
*/
@Slf4j
@Component
@@ -29,6 +37,7 @@ public class ToolInterceptor {
private final AgentEventPublisher eventPublisher;
private final ToolResultPresentationProcessor toolResultProcessor;
private final AgentProperties agentProperties;
private final ObjectProvider<List<McpSyncClient>> mcpSyncClientsProvider;
public List<ToolCallback> getWrappedTools(String sessionId, String messageId) {
return getWrappedTools(sessionId, messageId, AgentToolScope.ALL);
@@ -61,9 +70,16 @@ public class ToolInterceptor {
String toolName = getToolDefinition().name();
eventPublisher.toolRunning(sessionId, messageId, toolName);
try {
String result = tool.call(toolInput);
String result = executeWithRetry(tool, toolInput, toolName);
// Sanitize error stack traces — don't leak internals to LLM
// Unrecoverable raw system error responses (e.g. 404 HTML, connection reset)
if (result != null && isUnrecoverableErrorResponse(result)) {
log.error("Tool {} returned unrecoverable system error response: {}", toolName, result);
throw new UnrecoverableToolExecutionException(
"Tool " + toolName + " returned unrecoverable system error");
}
// Sanitize standard business/internal errors — don't leak stack traces to LLM
if (result != null && isErrorResponse(result)) {
log.warn("Tool {} returned error response, sanitizing", toolName);
return "{\"error\": \"Lỗi hệ thống khi gọi tool. Không thể xử lý yêu cầu. Vui lòng thử lại sau.\"}";
@@ -74,9 +90,12 @@ public class ToolInterceptor {
// Truncate if result is too large for the model's context window
return truncateResult(processed, toolName);
} catch (UnrecoverableToolExecutionException e) {
throw e;
} catch (Exception e) {
log.error("Tool {} threw exception: {}", toolName, e.getMessage(), e);
return "{\"error\": \"Lỗi hệ thống khi gọi tool. Không thể xử lý yêu cầu. Vui lòng thử lại sau.\"}";
log.error("Tool {} threw unrecoverable exception: {}", toolName, e.getMessage(), e);
throw new UnrecoverableToolExecutionException(
"Lỗi hệ thống không thể phục hồi khi gọi tool '" + toolName + "': " + e.getMessage(), e);
} finally {
eventPublisher.toolDone(sessionId, messageId, toolName);
}
@@ -84,6 +103,79 @@ public class ToolInterceptor {
};
}
private String executeWithRetry(ToolCallback tool, String toolInput, String toolName) throws Exception {
try {
return tool.call(toolInput);
} catch (Exception e) {
if (isConnectionOrSessionError(e)) {
log.warn("Detected MCP connection/session drop on tool {}: {}. Attempting auto-reconnect and retry...",
toolName, e.getMessage());
attemptReconnect();
try {
return tool.call(toolInput);
} catch (Exception retryEx) {
log.error("Retry after reconnect failed for tool {}: {}", toolName, retryEx.getMessage());
throw new UnrecoverableToolExecutionException("Không thể khôi phục kết nối MCP khi gọi tool " + toolName, retryEx);
}
}
throw e;
}
}
private void attemptReconnect() {
if (mcpSyncClientsProvider != null) {
List<McpSyncClient> clients = mcpSyncClientsProvider.getIfAvailable();
if (clients != null && !clients.isEmpty()) {
for (McpSyncClient client : clients) {
try {
log.info("Re-initializing MCP Sync Client session...");
client.initialize();
} catch (Exception ex) {
log.warn("Failed to re-initialize MCP Sync Client: {}", ex.getMessage());
}
}
}
}
}
public boolean isConnectionOrSessionError(Throwable t) {
if (t == null) {
return false;
}
Throwable current = t;
while (current != null) {
if (current instanceof EOFException
|| current instanceof ConnectException
|| current instanceof SocketException
|| current instanceof ClosedChannelException) {
return true;
}
String msg = current.getMessage();
if (msg != null) {
String lower = msg.toLowerCase();
if (lower.contains("404")
|| lower.contains("session not found")
|| lower.contains("session invalid")
|| lower.contains("connection reset")
|| lower.contains("connection refused")
|| lower.contains("stream closed")) {
return true;
}
}
current = current.getCause();
}
return false;
}
private boolean isUnrecoverableErrorResponse(String result) {
if (result == null) return false;
String lower = result.toLowerCase();
return lower.contains("404 session not found")
|| lower.contains("http status 404")
|| lower.contains("connection refused")
|| lower.contains("eofexception");
}
private boolean isErrorResponse(String result) {
return (result.contains("java.lang.") && result.contains("Exception") && result.contains("org.springframework."))
|| result.contains("HTTP Status 500 – Internal Server Error");

3
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/memory/ConversationKeyResolver.java
View File

@@ -6,6 +6,7 @@ import org.springframework.util.StringUtils;
/**
* Chooses the stable key used by memory and workflow state.
* Requires explicit conversationId from ChatRequest.
*/
@Component
public class ConversationKeyResolver {
@@ -14,6 +15,6 @@ public class ConversationKeyResolver {
if (request != null && StringUtils.hasText(request.conversationId())) {
return request.conversationId();
}
return connectionSessionId;
throw new IllegalArgumentException("conversationId is required in ChatRequest");
}
}

78
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/memory/InMemoryChatMemory.java
View File

@@ -1,78 +0,0 @@
package dev.sonpx.loyalty.agent.core.memory;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
/**
* In-memory chat memory with token-aware truncation.
*
* Improvements over original:
* - Configurable max tokens via AgentProperties
* - Vietnamese-aware token estimation (chars/2 instead of chars/4)
* - Preserves system messages during truncation
*/
@Component
public class InMemoryChatMemory implements ChatMemory {
private final Map<String, List<Message>> conversationHistory = new ConcurrentHashMap<>();
private final AgentProperties agentProperties;
public InMemoryChatMemory(AgentProperties agentProperties) {
this.agentProperties = agentProperties;
}
@Override
public void add(String conversationId, List<Message> messages) {
conversationHistory.compute(conversationId, (id, history) -> {
if (history == null) {
history = new ArrayList<>();
}
history.addAll(messages);
// Truncation: remove oldest non-system messages when over token limit
while (history.size() > 1 && estimateTokens(history) > agentProperties.getMaxMemoryTokens()) {
boolean removed = false;
for (int i = 0; i < history.size(); i++) {
if (!(history.get(i) instanceof org.springframework.ai.chat.messages.SystemMessage)) {
history.remove(i);
removed = true;
break;
}
}
// If only system messages remain, stop to avoid infinite loop
if (!removed) {
break;
}
}
return history;
});
}
/**
* Estimate token count for a list of messages.
* Uses configurable chars-per-token ratio (default 2 for Vietnamese text).
*/
private int estimateTokens(List<Message> messages) {
int chars = messages.stream()
.mapToInt(m -> m.getText() != null ? m.getText().length() : 0)
.sum();
return chars / agentProperties.getCharsPerToken();
}
@Override
public List<Message> get(String conversationId) {
return conversationHistory.getOrDefault(conversationId, new ArrayList<>());
}
@Override
public void clear(String conversationId) {
conversationHistory.remove(conversationId);
}
}

212
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/memory/JdbcChatMemory.java Normal file
View File

@@ -0,0 +1,212 @@
package dev.sonpx.loyalty.agent.core.memory;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
/**
* JDBC-backed ChatMemory for Spring AI with PostgreSQL persistence,
* Vietnamese token estimation, tool result truncation, and polymorphic message serialization.
*/
@Slf4j
@Component
public class JdbcChatMemory implements ChatMemory {
private static final int MAX_TOOL_RESULT_LENGTH = 4000;
private final JdbcTemplate jdbcTemplate;
private final AgentProperties agentProperties;
private final ObjectMapper objectMapper;
public JdbcChatMemory(JdbcTemplate jdbcTemplate, AgentProperties agentProperties, ObjectMapper objectMapper) {
this.jdbcTemplate = jdbcTemplate;
this.agentProperties = agentProperties;
this.objectMapper = objectMapper;
}
@Override
@Transactional
public void add(String conversationId, List<Message> messages) {
if (messages == null || messages.isEmpty()) {
return;
}
Long currentMaxSeq = jdbcTemplate.queryForObject(
"SELECT COALESCE(MAX(sequence_num), 0) FROM chat_memory_messages WHERE conversation_id = ?",
Long.class,
conversationId
);
long seq = (currentMaxSeq != null ? currentMaxSeq : 0);
for (Message msg : messages) {
seq++;
saveMessage(conversationId, seq, msg);
}
// Apply token-aware truncation
truncateIfExceedsTokenLimit(conversationId);
}
private void saveMessage(String conversationId, long sequenceNum, Message msg) {
String id = UUID.randomUUID().toString();
String messageType = msg.getMessageType().name();
String content = msg.getText();
String metadataJson = null;
if (msg instanceof AssistantMessage assistantMsg && assistantMsg.hasToolCalls()) {
try {
List<ToolCallDto> dtos = assistantMsg.getToolCalls().stream()
.map(tc -> new ToolCallDto(tc.id(), tc.type(), tc.name(), tc.arguments()))
.toList();
metadataJson = objectMapper.writeValueAsString(Map.of("toolCalls", dtos));
} catch (Exception e) {
log.error("Failed to serialize AssistantMessage toolCalls for conversation: {}", conversationId, e);
}
} else if (msg instanceof ToolResponseMessage toolMsg) {
try {
List<ToolResponseDto> dtos = toolMsg.getResponses().stream()
.map(tr -> {
String resData = tr.responseData();
if (resData != null && resData.length() > MAX_TOOL_RESULT_LENGTH) {
resData = resData.substring(0, MAX_TOOL_RESULT_LENGTH) + "\n... [truncated]";
}
return new ToolResponseDto(tr.id(), tr.name(), resData);
})
.toList();
metadataJson = objectMapper.writeValueAsString(Map.of("responses", dtos));
} catch (Exception e) {
log.error("Failed to serialize ToolResponseMessage responses for conversation: {}", conversationId, e);
}
}
jdbcTemplate.update(
"INSERT INTO chat_memory_messages (id, conversation_id, sequence_num, message_type, content, metadata, created_at) " +
"VALUES (?, ?, ?, ?, ?, ?::jsonb, ?)",
id, conversationId, sequenceNum, messageType, content, metadataJson, Timestamp.from(Instant.now())
);
}
private void truncateIfExceedsTokenLimit(String conversationId) {
List<MessageWrapper> history = loadMessageWrappers(conversationId);
int maxTokens = agentProperties.getMaxMemoryTokens();
int charsPerToken = agentProperties.getCharsPerToken();
while (history.size() > 1 && estimateTokensWrappers(history, charsPerToken) > maxTokens) {
int removeIndex = -1;
for (int i = 0; i < history.size(); i++) {
if (!(history.get(i).message instanceof SystemMessage)) {
removeIndex = i;
break;
}
}
if (removeIndex == -1) {
break; // Only system messages remain
}
MessageWrapper toRemove = history.remove(removeIndex);
jdbcTemplate.update("DELETE FROM chat_memory_messages WHERE id = ?", toRemove.id);
}
}
private int estimateTokensWrappers(List<MessageWrapper> history, int charsPerToken) {
int totalChars = history.stream()
.mapToInt(w -> w.message.getText() != null ? w.message.getText().length() : 0)
.sum();
return totalChars / (charsPerToken > 0 ? charsPerToken : 2);
}
@Override
public List<Message> get(String conversationId) {
List<MessageWrapper> wrappers = loadMessageWrappers(conversationId);
return wrappers.stream().map(w -> w.message).toList();
}
private List<MessageWrapper> loadMessageWrappers(String conversationId) {
String sql = "SELECT id, message_type, content, metadata FROM chat_memory_messages " +
"WHERE conversation_id = ? ORDER BY sequence_num ASC";
return jdbcTemplate.query(sql, (rs, rowNum) -> {
String id = rs.getString("id");
String messageTypeStr = rs.getString("message_type");
String content = rs.getString("content");
String metadataJson = rs.getString("metadata");
MessageType type = MessageType.valueOf(messageTypeStr);
Message message = deserializeMessage(type, content, metadataJson);
return new MessageWrapper(id, message);
}, conversationId);
}
private Message deserializeMessage(MessageType type, String content, String metadataJson) {
return switch (type) {
case USER -> new UserMessage(content != null ? content : "");
case SYSTEM -> new SystemMessage(content != null ? content : "");
case ASSISTANT -> {
if (metadataJson != null && !metadataJson.isBlank()) {
try {
Map<String, Object> map = objectMapper.readValue(metadataJson, new TypeReference<>() {});
if (map.containsKey("toolCalls")) {
List<ToolCallDto> dtos = objectMapper.convertValue(map.get("toolCalls"), new TypeReference<>() {});
List<AssistantMessage.ToolCall> toolCalls = dtos.stream()
.map(d -> new AssistantMessage.ToolCall(d.id(), d.type(), d.name(), d.arguments()))
.toList();
yield AssistantMessage.builder()
.content(content != null ? content : "")
.toolCalls(toolCalls)
.build();
}
} catch (Exception e) {
log.error("Error deserializing AssistantMessage metadata: {}", metadataJson, e);
}
}
yield new AssistantMessage(content != null ? content : "");
}
case TOOL -> {
if (metadataJson != null && !metadataJson.isBlank()) {
try {
Map<String, Object> map = objectMapper.readValue(metadataJson, new TypeReference<>() {});
if (map.containsKey("responses")) {
List<ToolResponseDto> dtos = objectMapper.convertValue(map.get("responses"), new TypeReference<>() {});
List<ToolResponseMessage.ToolResponse> responses = dtos.stream()
.map(d -> new ToolResponseMessage.ToolResponse(d.id(), d.name(), d.responseData()))
.toList();
yield ToolResponseMessage.builder()
.responses(responses)
.build();
}
} catch (Exception e) {
log.error("Error deserializing ToolResponseMessage metadata: {}", metadataJson, e);
}
}
yield ToolResponseMessage.builder().responses(List.of()).build();
}
};
}
@Override
@Transactional
public void clear(String conversationId) {
jdbcTemplate.update("DELETE FROM chat_memory_messages WHERE conversation_id = ?", conversationId);
}
private record MessageWrapper(String id, Message message) {}
private record ToolCallDto(String id, String type, String name, String arguments) {}
private record ToolResponseDto(String id, String name, String responseData) {}
}

3
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/orchestrator/PromptBuilder.java
View File

@@ -70,9 +70,6 @@ public class PromptBuilder {
+ request.conversationId()
+ ". You MUST provide this conversationId as a parameter to any tool that requires it.");
}
if (StringUtils.hasText(request.activeCampaignId())) {
context.add("Context: Active Campaign ID is " + request.activeCampaignId() + ".");
}
String prompt = request.prompt() == null ? "" : request.prompt();
if (context.isEmpty()) {

15
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/core/workflow/WorkflowExecutor.java
View File

@@ -43,6 +43,7 @@ public class WorkflowExecutor {
private final AgentEventPublisher eventPublisher;
private final AgentMemoryManager memoryManager;
private final WorkflowDataExtractor dataExtractor;
private final dev.sonpx.loyalty.agent.service.AgentPersistenceService persistenceService;
/** Active workflows per session */
private final Map<String, WorkflowState> activeWorkflows = new ConcurrentHashMap<>();
@@ -51,12 +52,14 @@ public class WorkflowExecutor {
ToolInterceptor toolInterceptor,
AgentEventPublisher eventPublisher,
AgentMemoryManager memoryManager,
WorkflowDataExtractor dataExtractor) {
WorkflowDataExtractor dataExtractor,
dev.sonpx.loyalty.agent.service.AgentPersistenceService persistenceService) {
this.chatClient = chatClient;
this.toolInterceptor = toolInterceptor;
this.eventPublisher = eventPublisher;
this.memoryManager = memoryManager;
this.dataExtractor = dataExtractor;
this.persistenceService = persistenceService;
}
/**
@@ -95,11 +98,6 @@ public class WorkflowExecutor {
log.info("Starting {} workflow for conversation: {}, connection: {}", intent, conversationKey, connectionSessionId);
// Inject activeCampaignId if available (useful for rule creation)
if (request.activeCampaignId() != null && !request.activeCampaignId().isBlank()) {
state.putData("Campaign ID", request.activeCampaignId());
}
// Build the initial prompt based on workflow type
String systemPrompt = getPhasePrompt(state);
String userPrompt = request.prompt() != null ? request.prompt() : "Bắt đầu tạo";
@@ -261,6 +259,9 @@ public class WorkflowExecutor {
String connectionSessionId, String conversationKey,
String messageId, boolean withTools,
AgentToolScope toolScope, Runnable onComplete, Runnable onError) {
persistenceService.saveUserMessage(conversationKey, messageId, userPrompt);
StringBuilder responseAccumulator = new StringBuilder();
var builder = chatClient.prompt()
.advisors(memoryManager.getAdvisor())
.advisors(a -> a
@@ -277,6 +278,7 @@ public class WorkflowExecutor {
builder.stream()
.content()
.doOnComplete(() -> {
persistenceService.saveAssistantMessage(conversationKey, responseAccumulator.toString());
if (onComplete != null) {
onComplete.run();
}
@@ -285,6 +287,7 @@ public class WorkflowExecutor {
.subscribe(
content -> {
if (!content.isEmpty()) {
responseAccumulator.append(content);
eventPublisher.token(connectionSessionId, messageId, content);
}
},

2
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/domain/ChatRequest.java
View File

@@ -1,4 +1,4 @@
package dev.sonpx.loyalty.agent.domain;
public record ChatRequest(String conversationId, String messageId, String prompt, String activeCampaignId) {
public record ChatRequest(String conversationId, String messageId, String prompt) {
}

14
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/domain/Conversation.java
View File

@@ -1,5 +1,9 @@
package dev.sonpx.loyalty.agent.domain;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import jakarta.persistence.Id;
import jakarta.persistence.Table;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
@@ -11,8 +15,18 @@ import java.time.Instant;
@Builder
@NoArgsConstructor
@AllArgsConstructor
@Entity
@Table(name = "conversations")
public class Conversation {
@Id
private String id;
@Column(name = "title")
private String title;
@Column(name = "created_at", nullable = false)
private Instant createdAt;
@Column(name = "updated_at", nullable = false)
private Instant updatedAt;
}

15
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/domain/Message.java
View File

@@ -1,5 +1,9 @@
package dev.sonpx.loyalty.agent.domain;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import jakarta.persistence.Id;
import jakarta.persistence.Table;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
@@ -11,10 +15,21 @@ import java.time.Instant;
@Builder
@NoArgsConstructor
@AllArgsConstructor
@Entity
@Table(name = "messages")
public class Message {
@Id
private String id;
@Column(name = "conversation_id", nullable = false)
private String conversationId;
@Column(name = "role", nullable = false)
private String role;
@Column(name = "content", nullable = false, columnDefinition = "TEXT")
private String content;
@Column(name = "created_at", nullable = false)
private Instant createdAt;
}

20
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/dto/ConversationSummaryDto.java Normal file
View File

@@ -0,0 +1,20 @@
package dev.sonpx.loyalty.agent.dto;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import java.time.Instant;
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class ConversationSummaryDto {
private String id;
private String title;
private String preview;
private Instant createdAt;
private Instant updatedAt;
}

17
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/exception/UnrecoverableToolExecutionException.java Normal file
View File

@@ -0,0 +1,17 @@
package dev.sonpx.loyalty.agent.exception;
/**
* Exception thrown when a tool execution encounters an unrecoverable system or connection error.
* Throwing this exception aborts the Spring AI ChatClient stream immediately,
* preventing small LLMs from entering an infinite ReAct retry loop.
*/
public class UnrecoverableToolExecutionException extends RuntimeException {
public UnrecoverableToolExecutionException(String message) {
super(message);
}
public UnrecoverableToolExecutionException(String message, Throwable cause) {
super(message, cause);
}
}

11
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/repository/ConversationRepository.java
View File

@@ -1,12 +1,13 @@
package dev.sonpx.loyalty.agent.repository;
import dev.sonpx.loyalty.agent.domain.Conversation;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.stereotype.Repository;
import java.util.Optional;
import java.util.List;
public interface ConversationRepository {
Conversation save(Conversation conversation);
Optional<Conversation> findById(String id);
List<Conversation> findAll();
@Repository
public interface ConversationRepository extends JpaRepository<Conversation, String> {
List<Conversation> findAllByOrderByUpdatedAtDesc();
}

31
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/repository/InMemoryConversationRepository.java
View File

@@ -1,31 +0,0 @@
package dev.sonpx.loyalty.agent.repository;
import dev.sonpx.loyalty.agent.domain.Conversation;
import org.springframework.stereotype.Repository;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
@Repository
public class InMemoryConversationRepository implements ConversationRepository {
private final Map<String, Conversation> store = new ConcurrentHashMap<>();
@Override
public Conversation save(Conversation conversation) {
store.put(conversation.getId(), conversation);
return conversation;
}
@Override
public Optional<Conversation> findById(String id) {
return Optional.ofNullable(store.get(id));
}
@Override
public List<Conversation> findAll() {
return new ArrayList<>(store.values());
}
}

30
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/repository/InMemoryMessageRepository.java
View File

@@ -1,30 +0,0 @@
package dev.sonpx.loyalty.agent.repository;
import dev.sonpx.loyalty.agent.domain.Message;
import org.springframework.stereotype.Repository;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
@Repository
public class InMemoryMessageRepository implements MessageRepository {
private final Map<String, Message> store = new ConcurrentHashMap<>();
@Override
public Message save(Message message) {
store.put(message.getId(), message);
return message;
}
@Override
public List<Message> findByConversationIdOrderByCreatedAtAsc(String conversationId) {
return store.values().stream()
.filter(m -> conversationId.equals(m.getConversationId()))
.sorted(Comparator.comparing(Message::getCreatedAt))
.collect(Collectors.toList());
}
}

6
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/repository/MessageRepository.java
View File

@@ -1,10 +1,12 @@
package dev.sonpx.loyalty.agent.repository;
import dev.sonpx.loyalty.agent.domain.Message;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.stereotype.Repository;
import java.util.List;
public interface MessageRepository {
Message save(Message message);
@Repository
public interface MessageRepository extends JpaRepository<Message, String> {
List<Message> findByConversationIdOrderByCreatedAtAsc(String conversationId);
}

6
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/service/AgentEventPublisher.java
View File

@@ -36,7 +36,11 @@ public class AgentEventPublisher {
}
private void send(String sessionId, AgentEvent event) {
log.debug("Publishing STOMP event type={} for session={}", event.type(), sessionId);
if ("TOKEN".equals(event.type())) {
log.trace("Publishing STOMP event type={} for session={}", event.type(), sessionId);
} else {
log.debug("Publishing STOMP event type={} for session={}", event.type(), sessionId);
}
messagingTemplate.convertAndSendToUser(sessionId, CHAT_EVENTS_DESTINATION, event);
}
}

109
loyalty-agent/src/main/java/dev/sonpx/loyalty/agent/service/AgentPersistenceService.java Normal file
View File

@@ -0,0 +1,109 @@
package dev.sonpx.loyalty.agent.service;
import dev.sonpx.loyalty.agent.domain.Conversation;
import dev.sonpx.loyalty.agent.domain.Message;
import dev.sonpx.loyalty.agent.repository.ConversationRepository;
import dev.sonpx.loyalty.agent.repository.MessageRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import java.time.Instant;
import java.util.UUID;
@Slf4j
@Service
public class AgentPersistenceService {
private final ConversationRepository conversationRepository;
private final MessageRepository messageRepository;
private final ChatClient chatChatClient;
public AgentPersistenceService(
ConversationRepository conversationRepository,
MessageRepository messageRepository,
@Qualifier("chatAgentChatClient") ChatClient chatChatClient) {
this.conversationRepository = conversationRepository;
this.messageRepository = messageRepository;
this.chatChatClient = chatChatClient;
}
public void saveUserMessage(String conversationId, String userMessageId, String content) {
try {
ensureConversationExists(conversationId, content);
Message message = Message.builder()
.id(userMessageId != null ? userMessageId : UUID.randomUUID().toString())
.conversationId(conversationId)
.role("user")
.content(content != null ? content : "")
.createdAt(Instant.now())
.build();
messageRepository.save(message);
} catch (Exception e) {
log.error("Failed to save user message for conversation: {}", conversationId, e);
}
}
public void saveAssistantMessage(String conversationId, String content) {
if (content == null || content.isBlank()) {
return;
}
try {
ensureConversationExists(conversationId, null);
Message message = Message.builder()
.id(UUID.randomUUID().toString())
.conversationId(conversationId)
.role("assistant")
.content(content)
.createdAt(Instant.now())
.build();
messageRepository.save(message);
} catch (Exception e) {
log.error("Failed to save assistant message for conversation: {}", conversationId, e);
}
}
private void ensureConversationExists(String conversationId, String userPromptForTitle) {
Conversation conversation = conversationRepository.findById(conversationId)
.orElseGet(() -> Conversation.builder()
.id(conversationId)
.title("Cuộc trò chuyện mới")
.createdAt(Instant.now())
.updatedAt(Instant.now())
.build());
conversation.setUpdatedAt(Instant.now());
if (userPromptForTitle != null && !userPromptForTitle.isBlank()) {
if (conversation.getTitle() == null || conversation.getTitle().isBlank() || "Cuộc trò chuyện mới".equalsIgnoreCase(conversation.getTitle())) {
String generatedTitle = generateTitleWithAgent(userPromptForTitle);
conversation.setTitle(generatedTitle);
}
}
conversationRepository.save(conversation);
}
private String generateTitleWithAgent(String content) {
try {
String prompt = "Hãy tạo 1 tiêu đề cực kỳ ngắn gọn (3 đến 6 từ), chính xác và súc tích bằng tiếng Việt đại diện cho chủ đề của tin nhắn sau. Không dùng dấu ngoặc kép, không giải thích dài dòng, chỉ trả về duy nhất tiêu đề:\n\n" + content;
String generatedTitle = chatChatClient.prompt()
.user(prompt)
.call()
.content();
if (generatedTitle != null && !generatedTitle.isBlank()) {
generatedTitle = generatedTitle.replaceAll("^[\"']|[\"']$", "").trim();
if (generatedTitle.length() > 60) {
generatedTitle = generatedTitle.substring(0, 60);
}
return generatedTitle;
}
} catch (Exception e) {
log.warn("Failed to generate title using Agent for content, falling back: {}", e.getMessage());
}
return content.length() > 40 ? content.substring(0, 40) + "..." : content;
}
}

21
loyalty-agent/src/main/resources/application.yml
View File

@@ -13,6 +13,23 @@ spring:
request-timeout: 3600000
application:
name: loyalty-agent-service
datasource:
url: ${SPRING_DATASOURCE_URL:jdbc:postgresql://192.168.99.242:5433/loyalty_agent}
username: ${SPRING_DATASOURCE_USERNAME:postgres}
password: ${SPRING_DATASOURCE_PASSWORD:Sonpx@1234}
driver-class-name: org.postgresql.Driver
jpa:
open-in-view: false
hibernate:
ddl-auto: validate
show-sql: false
properties:
hibernate:
format_sql: true
flyway:
enabled: true
baseline-on-migrate: true
locations: classpath:db/migration
ai:
ollama:
base-url: http://192.168.99.10:11434
@@ -35,6 +52,7 @@ logging:
pattern:
console: "%d{HH:mm:ss.SSS} %5p --- %-40.40logger{39} : %m%n"
level:
org.flywaydb: DEBUG
org.springframework.ai: INFO
dev.sonpx.loyalty: DEBUG
root: warn
@@ -56,6 +74,7 @@ agent:
creator:
base-url: ${CREATOR_OLLAMA_BASE_URL:http://192.168.99.10:11434}
model: ${CREATOR_MODEL:qwen3.5:4b}
temperature: ${CREATOR_MODEL_TEMPERATURE:0.2}
temperature: ${CREATOR_MODEL_TEMPERATURE:0.3}
num-ctx: ${CREATOR_MODEL_NUM_CTX:32768}
disable-thinking: ${CREATOR_MODEL_DISABLE_THINKING:false}

28
loyalty-agent/src/main/resources/db/migration/V1__init_agent_persistence.sql Normal file
View File

@@ -0,0 +1,28 @@
CREATE TABLE conversations (
id VARCHAR(255) PRIMARY KEY,
created_at TIMESTAMP WITH TIME ZONE NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE NOT NULL
);
CREATE TABLE messages (
id VARCHAR(255) PRIMARY KEY,
conversation_id VARCHAR(255) NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
role VARCHAR(50) NOT NULL,
content TEXT NOT NULL,
created_at TIMESTAMP WITH TIME ZONE NOT NULL
);
CREATE INDEX idx_messages_conversation_created ON messages(conversation_id, created_at);
CREATE TABLE chat_memory_messages (
id VARCHAR(255) PRIMARY KEY,
conversation_id VARCHAR(255) NOT NULL,
sequence_num BIGINT NOT NULL,
message_type VARCHAR(50) NOT NULL,
content TEXT,
metadata JSONB,
created_at TIMESTAMP WITH TIME ZONE NOT NULL,
CONSTRAINT uk_chat_memory_conversation_seq UNIQUE (conversation_id, sequence_num)
);
CREATE INDEX idx_chat_memory_conversation_created ON chat_memory_messages(conversation_id, created_at);

1
loyalty-agent/src/main/resources/db/migration/V2__add_title_to_conversations.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE conversations ADD COLUMN title VARCHAR(255);

124
loyalty-agent/src/test/java/dev/sonpx/loyalty/agent/core/executor/ToolInterceptorTest.java Normal file
View File

@@ -0,0 +1,124 @@
package dev.sonpx.loyalty.agent.core.executor;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import dev.sonpx.loyalty.agent.exception.UnrecoverableToolExecutionException;
import dev.sonpx.loyalty.agent.service.AgentEventPublisher;
import dev.sonpx.loyalty.agent.service.ToolResultPresentationProcessor;
import io.modelcontextprotocol.client.McpSyncClient;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.beans.factory.ObjectProvider;
import java.io.EOFException;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.*;
@ExtendWith(MockitoExtension.class)
public class ToolInterceptorTest {
@Mock
private ToolCallbackProvider toolCallbackProvider;
@Mock
private AgentEventPublisher eventPublisher;
@Mock
private ToolResultPresentationProcessor toolResultProcessor;
@Mock
private AgentProperties agentProperties;
@Mock
private ObjectProvider<List<McpSyncClient>> mcpSyncClientsProvider;
@Mock
private McpSyncClient mcpSyncClient;
@Mock
private ToolCallback mockTool;
@Mock
private ToolDefinition mockToolDefinition;
private ToolInterceptor toolInterceptor;
@BeforeEach
void setUp() {
lenient().when(mockToolDefinition.name()).thenReturn("test_tool");
lenient().when(mockTool.getToolDefinition()).thenReturn(mockToolDefinition);
lenient().when(toolCallbackProvider.getToolCallbacks()).thenReturn(new ToolCallback[]{mockTool});
toolInterceptor = new ToolInterceptor(
List.of(toolCallbackProvider),
eventPublisher,
toolResultProcessor,
agentProperties,
mcpSyncClientsProvider
);
}
@Test
void testSuccessfulToolExecution() {
when(mockTool.call("input")).thenReturn("{\"status\":\"ok\"}");
when(toolResultProcessor.process("{\"status\":\"ok\"}")).thenReturn("{\"status\":\"ok\"}");
when(agentProperties.getMaxToolResultChars()).thenReturn(1000);
List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");
assertThat(wrapped).hasSize(1);
String response = wrapped.get(0).call("input");
assertThat(response).isEqualTo("{\"status\":\"ok\"}");
verify(eventPublisher).toolRunning("session1", "msg1", "test_tool");
verify(eventPublisher).toolDone("session1", "msg1", "test_tool");
}
@Test
void testAutoReconnectAndRetryOnConnectionDrop() {
when(mcpSyncClientsProvider.getIfAvailable()).thenReturn(List.of(mcpSyncClient));
when(mockTool.call("input"))
.thenThrow(new RuntimeException("404 Session Not Found"))
.thenReturn("{\"recovered\":true}");
when(toolResultProcessor.process("{\"recovered\":true}")).thenReturn("{\"recovered\":true}");
when(agentProperties.getMaxToolResultChars()).thenReturn(1000);
List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");
String response = wrapped.get(0).call("input");
assertThat(response).isEqualTo("{\"recovered\":true}");
verify(mcpSyncClient, times(1)).initialize();
verify(mockTool, times(2)).call("input");
}
@Test
void testUnrecoverableErrorThrowsCircuitBreakerException() {
when(mcpSyncClientsProvider.getIfAvailable()).thenReturn(List.of(mcpSyncClient));
when(mockTool.call("input")).thenThrow(new RuntimeException(new EOFException("Connection reset by peer")));
List<ToolCallback> wrapped = toolInterceptor.getWrappedTools("session1", "msg1");
assertThatThrownBy(() -> wrapped.get(0).call("input"))
.isInstanceOf(UnrecoverableToolExecutionException.class)
.hasMessageContaining("Không thể khôi phục kết nối MCP khi gọi tool test_tool");
verify(eventPublisher).toolDone("session1", "msg1", "test_tool");
}
@Test
void testIsConnectionOrSessionErrorDetection() {
assertThat(toolInterceptor.isConnectionOrSessionError(new EOFException("EOF"))).isTrue();
assertThat(toolInterceptor.isConnectionOrSessionError(new RuntimeException("HTTP 404 Not Found"))).isTrue();
assertThat(toolInterceptor.isConnectionOrSessionError(new RuntimeException("Session not found"))).isTrue();
assertThat(toolInterceptor.isConnectionOrSessionError(new IllegalArgumentException("Invalid input"))).isFalse();
}
}

9
loyalty-agent/src/test/java/dev/sonpx/loyalty/agent/core/memory/ConversationKeyResolverTest.java
View File

@@ -1,6 +1,7 @@
package dev.sonpx.loyalty.agent.core.memory;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import dev.sonpx.loyalty.agent.domain.ChatRequest;
import org.junit.jupiter.api.Test;
@@ -11,15 +12,15 @@ class ConversationKeyResolverTest {
@Test
void usesConversationIdWhenPresent() {
ChatRequest request = new ChatRequest("conversation-1", "message-1", "hello", null);
ChatRequest request = new ChatRequest("conversation-1", "message-1", "hello");
assertEquals("conversation-1", resolver.resolve(request, "connection-1"));
}
@Test
void fallsBackToConnectionSessionId() {
ChatRequest request = new ChatRequest(" ", "message-1", "hello", null);
void throwsExceptionWhenConversationIdMissing() {
ChatRequest request = new ChatRequest(" ", "message-1", "hello");
assertEquals("connection-1", resolver.resolve(request, "connection-1"));
assertThrows(IllegalArgumentException.class, () -> resolver.resolve(request, "connection-1"));
}
}

182
loyalty-agent/src/test/java/dev/sonpx/loyalty/agent/core/memory/JdbcChatMemoryTest.java Normal file
View File

@@ -0,0 +1,182 @@
package dev.sonpx.loyalty.agent.core.memory;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import java.util.List;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;
@ExtendWith(MockitoExtension.class)
class JdbcChatMemoryTest {
@Mock
private JdbcTemplate jdbcTemplate;
@Mock
private AgentProperties agentProperties;
private final ObjectMapper objectMapper = new ObjectMapper();
private JdbcChatMemory jdbcChatMemory;
@BeforeEach
void setUp() {
lenient().when(agentProperties.getMaxMemoryTokens()).thenReturn(1000);
lenient().when(agentProperties.getCharsPerToken()).thenReturn(2);
jdbcChatMemory = new JdbcChatMemory(jdbcTemplate, agentProperties, objectMapper);
}
@Test
@DisplayName("Should save UserMessage and SystemMessage to database")
void testAddUserAndSystemMessage() {
when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-1"))).thenReturn(0L);
List<Message> messages = List.of(
new SystemMessage("System prompt"),
new UserMessage("Hello AI")
);
jdbcChatMemory.add("conv-1", messages);
verify(jdbcTemplate, times(2)).update(
contains("INSERT INTO chat_memory_messages"),
anyString(), eq("conv-1"), anyLong(), anyString(), any(), any(), any()
);
}
@Test
@DisplayName("Should serialize AssistantMessage tool calls into metadata JSON")
void testAddAssistantMessageWithToolCalls() {
when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-2"))).thenReturn(0L);
AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall("call-123", "function", "get_campaign", "{\"id\":1}");
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("Checking campaign...")
.toolCalls(List.of(toolCall))
.build();
jdbcChatMemory.add("conv-2", List.of(assistantMessage));
ArgumentCaptor<String> metadataCaptor = ArgumentCaptor.forClass(String.class);
verify(jdbcTemplate).update(
contains("INSERT INTO chat_memory_messages"),
anyString(), eq("conv-2"), eq(1L), eq("ASSISTANT"), eq("Checking campaign..."), metadataCaptor.capture(), any()
);
String metadataJson = metadataCaptor.getValue();
assertNotNull(metadataJson);
assertTrue(metadataJson.contains("call-123"));
assertTrue(metadataJson.contains("get_campaign"));
}
@Test
@DisplayName("Should truncate ToolResponseMessage when response data exceeds 4000 chars")
void testAddToolResponseMessageTruncation() {
when(jdbcTemplate.queryForObject(anyString(), eq(Long.class), eq("conv-3"))).thenReturn(0L);
String longResponseData = "A".repeat(5000);
ToolResponseMessage.ToolResponse toolResponse = new ToolResponseMessage.ToolResponse("call-123", "get_campaign", longResponseData);
ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
.responses(List.of(toolResponse))
.build();
jdbcChatMemory.add("conv-3", List.of(toolResponseMessage));
ArgumentCaptor<String> metadataCaptor = ArgumentCaptor.forClass(String.class);
verify(jdbcTemplate).update(
contains("INSERT INTO chat_memory_messages"),
anyString(), eq("conv-3"), eq(1L), eq("TOOL"), eq(""), metadataCaptor.capture(), any()
);
String metadataJson = metadataCaptor.getValue();
assertNotNull(metadataJson);
assertTrue(metadataJson.contains("... [truncated]"));
assertFalse(metadataJson.contains(longResponseData));
}
@Test
@DisplayName("Should deserialize AssistantMessage with tool calls correctly on get()")
@SuppressWarnings("unchecked")
void testGetAssistantMessageWithToolCalls() {
String metadataJson = "{\"toolCalls\":[{\"id\":\"call-99\",\"type\":\"function\",\"name\":\"check_rule\",\"arguments\":\"{\\\"ruleId\\\":5}\"}]}";
when(jdbcTemplate.query(anyString(), any(RowMapper.class), eq("conv-4")))
.thenAnswer(invocation -> {
RowMapper<?> mapper = invocation.getArgument(1);
java.sql.ResultSet rs = mock(java.sql.ResultSet.class);
when(rs.getString("id")).thenReturn("msg-1");
when(rs.getString("message_type")).thenReturn("ASSISTANT");
when(rs.getString("content")).thenReturn("Found rule");
when(rs.getString("metadata")).thenReturn(metadataJson);
Object wrapper = mapper.mapRow(rs, 0);
return List.of(wrapper);
});
List<Message> messages = jdbcChatMemory.get("conv-4");
assertEquals(1, messages.size());
assertTrue(messages.get(0) instanceof AssistantMessage);
AssistantMessage assistantMsg = (AssistantMessage) messages.get(0);
assertEquals("Found rule", assistantMsg.getText());
assertTrue(assistantMsg.hasToolCalls());
assertEquals(1, assistantMsg.getToolCalls().size());
assertEquals("call-99", assistantMsg.getToolCalls().get(0).id());
assertEquals("check_rule", assistantMsg.getToolCalls().get(0).name());
}
@Test
@DisplayName("Should deserialize ToolResponseMessage with responses correctly on get()")
@SuppressWarnings("unchecked")
void testGetToolResponseMessage() {
String metadataJson = "{\"responses\":[{\"id\":\"call-99\",\"name\":\"check_rule\",\"responseData\":\"{\\\"success\\\":true}\"}]}";
when(jdbcTemplate.query(anyString(), any(RowMapper.class), eq("conv-5")))
.thenAnswer(invocation -> {
RowMapper<?> mapper = invocation.getArgument(1);
java.sql.ResultSet rs = mock(java.sql.ResultSet.class);
when(rs.getString("id")).thenReturn("msg-2");
when(rs.getString("message_type")).thenReturn("TOOL");
when(rs.getString("content")).thenReturn("");
when(rs.getString("metadata")).thenReturn(metadataJson);
Object wrapper = mapper.mapRow(rs, 0);
return List.of(wrapper);
});
List<Message> messages = jdbcChatMemory.get("conv-5");
assertEquals(1, messages.size());
assertTrue(messages.get(0) instanceof ToolResponseMessage);
ToolResponseMessage toolMsg = (ToolResponseMessage) messages.get(0);
assertEquals(1, toolMsg.getResponses().size());
assertEquals("call-99", toolMsg.getResponses().get(0).id());
assertEquals("{\"success\":true}", toolMsg.getResponses().get(0).responseData());
}
@Test
@DisplayName("Should clear chat memory messages for a conversation")
void testClear() {
jdbcChatMemory.clear("conv-6");
verify(jdbcTemplate).update("DELETE FROM chat_memory_messages WHERE conversation_id = ?", "conv-6");
}
}