fix(ai): correct conversation message handling (#4208)
Signed-off-by: yuluo-yx <yuluo08290126@gmail.com>
diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java
index 94c7605..b642d62 100644
--- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java
+++ b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java
@@ -34,7 +34,6 @@
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
-import org.springframework.util.CollectionUtils;
import reactor.core.publisher.Flux;
import java.util.Collections;
@@ -99,8 +98,7 @@
ChatRequestContext context = ChatRequestContext.builder()
.message(message)
.conversationId(conversationId)
- .conversationHistory(CollectionUtils.isEmpty(conversation.getMessages()) ? null
- : conversation.getMessages().subList(0, conversation.getMessages().size() - 1))
+ .conversationHistory(messages)
.build();
// Stream response from AI service
diff --git a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java b/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java
new file mode 100644
index 0000000..f46c5dc
--- /dev/null
+++ b/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.ai.service.impl;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicLong;
+import org.apache.hertzbeat.ai.dao.ChatConversationDao;
+import org.apache.hertzbeat.ai.dao.ChatMessageDao;
+import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
+import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
+import org.apache.hertzbeat.ai.service.ChatClientProviderService;
+import org.apache.hertzbeat.common.entity.ai.ChatConversation;
+import org.apache.hertzbeat.common.entity.ai.ChatMessage;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.InjectMocks;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.http.codec.ServerSentEvent;
+import reactor.core.publisher.Flux;
+
+/**
+ * Tests multi-turn conversation context handling in {@link ConversationServiceImpl}.
+ */
+@ExtendWith(MockitoExtension.class)
+class ConversationServiceImplTest {
+
+ private static final long CONVERSATION_ID = 1L;
+
+ @Mock
+ private ChatConversationDao conversationDao;
+
+ @Mock
+ private ChatMessageDao messageDao;
+
+ @Mock
+ private ChatClientProviderService chatClientProviderService;
+
+ @InjectMocks
+ private ConversationServiceImpl conversationService;
+
+ @Test
+ void streamChatShouldKeepCompleteConversationHistory() {
+ ChatConversation conversation = ChatConversation.builder()
+ .id(CONVERSATION_ID)
+ .title("已命名会话")
+ .build();
+ List<ChatMessage> history = List.of(
+ ChatMessage.builder()
+ .id(11L)
+ .conversationId(CONVERSATION_ID)
+ .role("user")
+ .content("上一轮问题")
+ .build(),
+ ChatMessage.builder()
+ .id(12L)
+ .conversationId(CONVERSATION_ID)
+ .role("assistant")
+ .content("上一轮回答")
+ .build());
+ AtomicLong messageId = new AtomicLong(20L);
+
+ when(chatClientProviderService.isConfigured()).thenReturn(true);
+ when(conversationDao.findById(CONVERSATION_ID)).thenReturn(Optional.of(conversation));
+ when(messageDao.findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID)).thenReturn(history);
+ when(messageDao.save(any(ChatMessage.class))).thenAnswer(invocation -> {
+ ChatMessage savedMessage = invocation.getArgument(0);
+ savedMessage.setId(messageId.getAndIncrement());
+ return savedMessage;
+ });
+ when(chatClientProviderService.streamChat(any(ChatRequestContext.class)))
+ .thenReturn(Flux.just("本轮回答"));
+
+ List<ServerSentEvent<ChatResponseChunk>> events = conversationService
+ .streamChat("本轮问题", CONVERSATION_ID)
+ .collectList()
+ .block();
+
+ assertNotNull(events);
+ assertEquals(2, events.size());
+ ArgumentCaptor<ChatRequestContext> contextCaptor = ArgumentCaptor.forClass(ChatRequestContext.class);
+ verify(chatClientProviderService).streamChat(contextCaptor.capture());
+ assertEquals(history, contextCaptor.getValue().getConversationHistory());
+ }
+}
diff --git a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/ai/ChatMessage.java b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/ai/ChatMessage.java
index 3e7568b..8938ba4 100644
--- a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/ai/ChatMessage.java
+++ b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/ai/ChatMessage.java
@@ -18,6 +18,7 @@
package org.apache.hertzbeat.common.entity.ai;
import static io.swagger.v3.oas.annotations.media.Schema.AccessMode.READ_ONLY;
+import com.fasterxml.jackson.annotation.JsonIgnore;
import io.swagger.v3.oas.annotations.media.Schema;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
@@ -61,12 +62,13 @@
private Long id;
@Schema(title = "conversation id")
- @Column(name = "conversation_id", insertable = false, updatable = false)
+ @Column(name = "conversation_id")
private Long conversationId;
+ @JsonIgnore
@Schema(title = "conversation", hidden = true)
@ManyToOne
- @JoinColumn(name = "conversation_id")
+ @JoinColumn(name = "conversation_id", insertable = false, updatable = false)
private ChatConversation conversation;
@Schema(title = "message content")
diff --git a/hertzbeat-common-spring/src/test/java/org/apache/hertzbeat/common/entity/ai/ChatMessageTest.java b/hertzbeat-common-spring/src/test/java/org/apache/hertzbeat/common/entity/ai/ChatMessageTest.java
new file mode 100644
index 0000000..801e9ab
--- /dev/null
+++ b/hertzbeat-common-spring/src/test/java/org/apache/hertzbeat/common/entity/ai/ChatMessageTest.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.common.entity.ai;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import java.util.List;
+import org.apache.hertzbeat.common.util.JsonUtil;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Tests AI conversation message serialization.
+ */
+class ChatMessageTest {
+
+ @Test
+ void serializationShouldNotRecurseThroughConversation() {
+ ChatConversation conversation = ChatConversation.builder()
+ .id(1L)
+ .title("序列化测试会话")
+ .build();
+ ChatMessage message = ChatMessage.builder()
+ .id(2L)
+ .conversationId(conversation.getId())
+ .conversation(conversation)
+ .role("assistant")
+ .content("序列化测试消息")
+ .build();
+ conversation.setMessages(List.of(message));
+
+ String json = JsonUtil.toJson(conversation);
+
+ assertNotNull(json);
+ assertFalse(json.contains("\"conversation\""));
+ }
+}
diff --git a/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/dao/ChatMessageDaoTest.java b/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/dao/ChatMessageDaoTest.java
new file mode 100644
index 0000000..bbdc3cb
--- /dev/null
+++ b/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/dao/ChatMessageDaoTest.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.startup.dao;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import jakarta.annotation.Resource;
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.PersistenceContext;
+import java.util.List;
+import org.apache.hertzbeat.ai.dao.ChatConversationDao;
+import org.apache.hertzbeat.ai.dao.ChatMessageDao;
+import org.apache.hertzbeat.common.entity.ai.ChatConversation;
+import org.apache.hertzbeat.common.entity.ai.ChatMessage;
+import org.apache.hertzbeat.startup.AbstractSpringIntegrationTest;
+import org.junit.jupiter.api.Test;
+import org.springframework.transaction.annotation.Transactional;
+
+/**
+ * Tests persistence mapping for AI conversation messages.
+ */
+@Transactional
+class ChatMessageDaoTest extends AbstractSpringIntegrationTest {
+
+ @Resource
+ private ChatConversationDao conversationDao;
+
+ @Resource
+ private ChatMessageDao messageDao;
+
+ @PersistenceContext
+ private EntityManager entityManager;
+
+ @Test
+ void saveMessageShouldPersistConversationId() {
+ ChatConversation conversation = conversationDao.saveAndFlush(
+ ChatConversation.builder().title("映射测试会话").build());
+ messageDao.saveAndFlush(ChatMessage.builder()
+ .conversationId(conversation.getId())
+ .role("user")
+ .content("映射测试消息")
+ .build());
+ entityManager.clear();
+
+ List<ChatMessage> messages = messageDao
+ .findByConversationIdOrderByGmtCreateAsc(conversation.getId());
+
+ assertEquals(1, messages.size());
+ assertEquals(conversation.getId(), messages.getFirst().getConversationId());
+ }
+}
