Skip to content

Commit 2bb4f43

Browse files
LiuTianyoulynx009Aias00yuluo-yxDuansg
authored
[improve] optimize the function of creating monitoring using AI (#3922)
Signed-off-by: aias00 <liuhongyu@apache.org> Co-authored-by: lynx009 <2030509072@qq.com> Co-authored-by: aias00 <liuhongyu@apache.org> Co-authored-by: shown <yuluo08290126@gmail.com> Co-authored-by: Duansg <siguoduan@gmail.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: manshi <lty496227@digital-engine.com>
1 parent 695be42 commit 2bb4f43

22 files changed

Lines changed: 791 additions & 227 deletions

File tree

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/controller/ChatController.java‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.hertzbeat.ai.config.McpContextHolder;
2727
import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
2828
import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
29+
import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
2930
import org.apache.hertzbeat.ai.service.ConversationService;
3031
import org.apache.hertzbeat.common.entity.ai.ChatConversation;
3132
import org.apache.hertzbeat.common.entity.dto.Message;
@@ -78,12 +79,12 @@ public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(@Valid @RequestBody C
7879
McpContextHolder.setSubject(subject);
7980
if (context.getMessage() == null || context.getMessage().trim().isEmpty()) {
8081
ChatResponseChunk errorResponse = ChatResponseChunk.builder()
81-
.conversationId(context.getConversationId())
82-
.response("Error: Message cannot be empty")
83-
.build();
82+
.conversationId(context.getConversationId())
83+
.response("Error: Message cannot be empty")
84+
.build();
8485
return Flux.just(ServerSentEvent.builder(errorResponse)
85-
.event("error")
86-
.build());
86+
.event("error")
87+
.build());
8788
}
8889

8990
log.info("Received streaming chat request for conversation: {}", context.getConversationId());
@@ -92,12 +93,12 @@ public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(@Valid @RequestBody C
9293
} catch (Exception e) {
9394
log.error("Error in stream chat endpoint: ", e);
9495
ChatResponseChunk errorResponse = ChatResponseChunk.builder()
95-
.conversationId(context.getConversationId())
96-
.response("An error occurred: " + e.getMessage())
97-
.build();
96+
.conversationId(context.getConversationId())
97+
.response("An error occurred: " + e.getMessage())
98+
.build();
9899
return Flux.just(ServerSentEvent.builder(errorResponse)
99-
.event("error")
100-
.build());
100+
.event("error")
101+
.build());
101102
}
102103
}
103104

@@ -134,7 +135,7 @@ public ResponseEntity<Message<List<ChatConversation>>> listConversations() {
134135
@GetMapping(path = "/conversations/{conversationId}")
135136
@Operation(summary = "Get conversation history", description = "Get detailed information and message history for a specific conversation")
136137
public ResponseEntity<Message<ChatConversation>> getConversation(
137-
@Parameter(description = "Conversation ID", example = "12345678") @PathVariable(value = "conversationId") Long conversationId) {
138+
@Parameter(description = "Conversation ID", example = "12345678") @PathVariable(value = "conversationId") Long conversationId) {
138139
ChatConversation conversation = conversationService.getConversation(conversationId);
139140
return ResponseEntity.ok(Message.success(conversation));
140141
}
@@ -148,8 +149,21 @@ public ResponseEntity<Message<ChatConversation>> getConversation(
148149
@DeleteMapping(path = "/conversations/{conversationId}")
149150
@Operation(summary = "Delete conversation", description = "Delete a specific conversation and all its messages")
150151
public ResponseEntity<Message<Void>> deleteConversation(
151-
@Parameter(description = "Conversation ID", example = "2345678") @PathVariable("conversationId") Long conversationId) {
152+
@Parameter(description = "Conversation ID", example = "2345678") @PathVariable("conversationId") Long conversationId) {
152153
conversationService.deleteConversation(conversationId);
153154
return ResponseEntity.ok(Message.success());
154155
}
156+
157+
/**
158+
* Save data submitted by secure form
159+
* @param securityData security data
160+
* @return save result
161+
*/
162+
@PostMapping(path = "/security")
163+
@Operation(summary = "save security data", description = "Save security data")
164+
public ResponseEntity<Message<Boolean>> commitSecurityData(@Valid @RequestBody SecurityData securityData) {
165+
return ResponseEntity.ok(Message.success(conversationService.saveSecurityData(securityData)));
166+
}
167+
168+
155169
}

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/pojo/dto/SecurityData.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.hertzbeat.ai.pojo.dto;
19+
20+
import io.swagger.v3.oas.annotations.media.Schema;
21+
import jakarta.validation.constraints.NotNull;
22+
import lombok.Data;
23+
24+
/**
25+
* security data
26+
*/
27+
@Data
28+
public class SecurityData {
29+
30+
@NotNull
31+
@Schema(description = "Conversation ID", example = "123")
32+
private Long conversationId;
33+
34+
@NotNull
35+
@Schema(description = "security data", example = "{\"password\":\"xxxxx\"}")
36+
private String securityData;
37+
38+
}

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/ConversationService.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.hertzbeat.ai.service;
2020

2121
import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
22+
import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
2223
import org.apache.hertzbeat.common.entity.ai.ChatConversation;
2324
import org.springframework.http.codec.ServerSentEvent;
2425
import reactor.core.publisher.Flux;
@@ -33,7 +34,7 @@ public interface ConversationService {
3334
/**
3435
* Send a message and receive a streaming response
3536
*
36-
* @param message The user's message
37+
* @param message The user's message
3738
* @param conversationId Optional conversation ID for continuing a chat
3839
* @return Flux of ServerSentEvent for streaming the response
3940
*/
@@ -67,4 +68,13 @@ public interface ConversationService {
6768
* @param conversationId Conversation ID to delete
6869
*/
6970
void deleteConversation(Long conversationId);
71+
72+
/**
73+
* save security data for a conversation
74+
*
75+
* @param securityData securityData
76+
* @return save result
77+
*/
78+
Boolean saveSecurityData(SecurityData securityData);
79+
7080
}

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImpl.java‎

Lines changed: 59 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@
1818

1919
package org.apache.hertzbeat.ai.service.impl;
2020

21+
import java.nio.charset.StandardCharsets;
22+
import java.util.HashMap;
23+
import java.util.Map;
24+
import java.util.Objects;
2125
import lombok.extern.slf4j.Slf4j;
2226
import org.apache.hertzbeat.ai.sop.model.SopDefinition;
2327
import org.apache.hertzbeat.ai.sop.model.SopParameter;
@@ -27,9 +31,12 @@
2731
import org.apache.hertzbeat.ai.service.ChatClientProviderService;
2832
import org.apache.hertzbeat.base.dao.GeneralConfigDao;
2933
import org.apache.hertzbeat.common.entity.manager.GeneralConfig;
34+
import org.apache.hertzbeat.common.support.event.AiProviderConfigChangeEvent;
3035
import org.apache.hertzbeat.common.util.JsonUtil;
36+
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
3137
import org.springframework.beans.factory.annotation.Value;
3238
import org.springframework.context.annotation.Lazy;
39+
import org.springframework.context.event.EventListener;
3340
import org.springframework.core.io.Resource;
3441
import org.springframework.stereotype.Service;
3542
import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
@@ -44,14 +51,12 @@
4451
import reactor.core.publisher.Flux;
4552

4653
import java.io.IOException;
47-
import java.nio.charset.StandardCharsets;
4854
import java.util.ArrayList;
4955
import java.util.List;
5056

5157
/**
52-
* Implementation of the {@link ChatClientProviderService}.
53-
* Provides functionality to interact with the ChatClient for handling chat
54-
* messages.
58+
* Implementation of the {@link ChatClientProviderService}. Provides functionality to interact with the ChatClient for
59+
* handling chat messages.
5560
*/
5661
@Slf4j
5762
@Service
@@ -64,21 +69,28 @@ public class ChatClientProviderServiceImpl implements ChatClientProviderService
6469

6570
private final GeneralConfigDao generalConfigDao;
6671

72+
private ModelProviderConfig modelProviderConfig;
73+
74+
6775
private final SkillRegistry skillRegistry;
68-
76+
6977
@Autowired
7078
@Qualifier("hertzbeatTools")
7179
private ToolCallbackProvider toolCallbackProvider;
72-
80+
7381
private boolean isConfigured = false;
7482

7583
@Value("classpath:/prompt/system-message.st")
7684
private Resource systemResource;
7785

86+
@Value("classpath:/prompt/extra-message-protected.st")
87+
private Resource extraResourceProtected;
88+
89+
7890
@Autowired
79-
public ChatClientProviderServiceImpl(ApplicationContext applicationContext,
80-
GeneralConfigDao generalConfigDao,
81-
@Lazy SkillRegistry skillRegistry) {
91+
public ChatClientProviderServiceImpl(ApplicationContext applicationContext,
92+
GeneralConfigDao generalConfigDao,
93+
@Lazy SkillRegistry skillRegistry) {
8294
this.applicationContext = applicationContext;
8395
this.generalConfigDao = generalConfigDao;
8496
this.skillRegistry = skillRegistry;
@@ -89,7 +101,7 @@ public Flux<String> streamChat(ChatRequestContext context) {
89101
try {
90102
// Get the current (potentially refreshed) ChatClient instance
91103
ChatClient chatClient = applicationContext.getBean("openAiChatClient", ChatClient.class);
92-
104+
93105
List<Message> messages = new ArrayList<>();
94106

95107
// Add conversation history if available
@@ -112,13 +124,13 @@ public Flux<String> streamChat(ChatRequestContext context) {
112124
String systemPrompt = buildSystemPrompt(context.getConversationId());
113125

114126
return chatClient.prompt()
115-
.messages(messages)
116-
.system(systemPrompt)
117-
.toolCallbacks(toolCallbackProvider)
118-
.stream()
119-
.content()
120-
.doOnComplete(() -> log.info("Streaming completed for conversation: {}", context.getConversationId()))
121-
.doOnError(error -> log.error("Error in streaming chat: {}", error.getMessage(), error));
127+
.messages(messages)
128+
.system(systemPrompt)
129+
.toolCallbacks(toolCallbackProvider)
130+
.stream()
131+
.content()
132+
.doOnComplete(() -> log.info("Streaming completed for conversation: {}", context.getConversationId()))
133+
.doOnError(error -> log.error("Error in streaming chat: {}", error.getMessage(), error));
122134

123135
} catch (Exception e) {
124136
log.error("Error setting up streaming chat: {}", e.getMessage(), e);
@@ -133,30 +145,43 @@ private String buildSystemPrompt(Long conversationId) {
133145
try {
134146
String template = systemResource.getContentAsString(StandardCharsets.UTF_8);
135147
String skillsList = generateSkillsList();
136-
return template
137-
.replace(SKILLS_PLACEHOLDER, skillsList)
138-
.replace(CONVERSATION_ID_PLACEHOLDER, String.valueOf(conversationId));
148+
template = template
149+
.replace(SKILLS_PLACEHOLDER, skillsList)
150+
.replace(CONVERSATION_ID_PLACEHOLDER, String.valueOf(conversationId));
151+
152+
// add extra prompt for protected model to guide it to use protected tools
153+
if (Objects.equals(modelProviderConfig.getParticipationModel(), "PROTECTED")) {
154+
Map<String, Object> metadata = new HashMap<>();
155+
metadata.put("conversationId", conversationId);
156+
return template + SystemPromptTemplate.builder().resource(extraResourceProtected).build()
157+
.create(metadata)
158+
.getContents();
159+
} else {
160+
return template;
161+
}
162+
139163
} catch (IOException e) {
140164
log.error("Failed to read system prompt template: {}", e.getMessage());
141165
return "";
142166
}
143167
}
144168

169+
145170
/**
146171
* Generate a formatted list of available skills for the system prompt.
147172
*/
148173
private String generateSkillsList() {
149174
List<SopDefinition> skills = skillRegistry.getAllSkills();
150-
175+
151176
if (skills.isEmpty()) {
152177
return "No skills currently available. Use listSkills tool to refresh.";
153178
}
154-
179+
155180
StringBuilder sb = new StringBuilder();
156181
for (SopDefinition skill : skills) {
157182
sb.append("- **").append(skill.getName()).append("**: ");
158183
sb.append(skill.getDescription());
159-
184+
160185
// Add parameter hints
161186
if (skill.getParameters() != null && !skill.getParameters().isEmpty()) {
162187
sb.append(" (requires: ");
@@ -171,16 +196,24 @@ private String generateSkillsList() {
171196
}
172197
sb.append("\n");
173198
}
174-
199+
175200
return sb.toString();
176201
}
177202

203+
@EventListener(AiProviderConfigChangeEvent.class)
204+
public void onAiProviderConfigChange(AiProviderConfigChangeEvent event) {
205+
GeneralConfig providerConfig = generalConfigDao.findByType("provider");
206+
this.modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
207+
}
208+
178209
@Override
179210
public boolean isConfigured() {
180211
if (!isConfigured) {
181212
GeneralConfig providerConfig = generalConfigDao.findByType("provider");
182-
ModelProviderConfig modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
183-
isConfigured = modelProviderConfig != null && modelProviderConfig.getApiKey() != null;
213+
ModelProviderConfig modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(),
214+
ModelProviderConfig.class);
215+
isConfigured = modelProviderConfig != null && modelProviderConfig.getApiKey() != null;
216+
this.modelProviderConfig = modelProviderConfig;
184217
}
185218
return isConfigured;
186219
}

0 commit comments

Comments
 (0)