[improve] optimize the function of creating monitoring using AI

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

@@ -26,6 +26,7 @@
 import org.apache.hertzbeat.ai.config.McpContextHolder;
 import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
 import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
+import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
 import org.apache.hertzbeat.ai.service.ConversationService;
 import org.apache.hertzbeat.common.entity.ai.ChatConversation;
 import org.apache.hertzbeat.common.entity.dto.Message;
@@ -78,12 +79,12 @@ public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(@Valid @RequestBody C
             McpContextHolder.setSubject(subject);
             if (context.getMessage() == null || context.getMessage().trim().isEmpty()) {
                 ChatResponseChunk errorResponse = ChatResponseChunk.builder()
-                        .conversationId(context.getConversationId())
-                        .response("Error: Message cannot be empty")
-                        .build();
+                    .conversationId(context.getConversationId())
+                    .response("Error: Message cannot be empty")
+                    .build();
                 return Flux.just(ServerSentEvent.builder(errorResponse)
-                        .event("error")
-                        .build());
+                    .event("error")
+                    .build());
             }
 
             log.info("Received streaming chat request for conversation: {}", context.getConversationId());
@@ -92,12 +93,12 @@ public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(@Valid @RequestBody C
         } catch (Exception e) {
             log.error("Error in stream chat endpoint: ", e);
             ChatResponseChunk errorResponse = ChatResponseChunk.builder()
-                    .conversationId(context.getConversationId())
-                    .response("An error occurred: " + e.getMessage())
-                    .build();
+                .conversationId(context.getConversationId())
+                .response("An error occurred: " + e.getMessage())
+                .build();
             return Flux.just(ServerSentEvent.builder(errorResponse)
-                    .event("error")
-                    .build());
+                .event("error")
+                .build());
         }
     }
 
@@ -134,7 +135,7 @@ public ResponseEntity<Message<List<ChatConversation>>> listConversations() {
     @GetMapping(path = "/conversations/{conversationId}")
     @Operation(summary = "Get conversation history", description = "Get detailed information and message history for a specific conversation")
     public ResponseEntity<Message<ChatConversation>> getConversation(
-            @Parameter(description = "Conversation ID", example = "12345678") @PathVariable(value = "conversationId") Long conversationId) {
+        @Parameter(description = "Conversation ID", example = "12345678") @PathVariable(value = "conversationId") Long conversationId) {
         ChatConversation conversation = conversationService.getConversation(conversationId);
         return ResponseEntity.ok(Message.success(conversation));
     }
@@ -148,8 +149,21 @@ public ResponseEntity<Message<ChatConversation>> getConversation(
     @DeleteMapping(path = "/conversations/{conversationId}")
     @Operation(summary = "Delete conversation", description = "Delete a specific conversation and all its messages")
     public ResponseEntity<Message<Void>> deleteConversation(
-            @Parameter(description = "Conversation ID", example = "2345678") @PathVariable("conversationId") Long conversationId) {
+        @Parameter(description = "Conversation ID", example = "2345678") @PathVariable("conversationId") Long conversationId) {
         conversationService.deleteConversation(conversationId);
         return ResponseEntity.ok(Message.success());
     }
+
+    /**
+     * Save data submitted by secure form
+     * @param securityData security data
+     * @return save result
+     */
+    @PostMapping(path = "/security")
+    @Operation(summary = "save security data", description = "Save security data")
+    public ResponseEntity<Message<Boolean>> commitSecurityData(@Valid @RequestBody SecurityData securityData) {
+        return ResponseEntity.ok(Message.success(conversationService.saveSecurityData(securityData)));
+    }
+
+
 }

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

@@ -0,0 +1,38 @@
+/*
+ * 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.pojo.dto;
+
+import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.NotNull;
+import lombok.Data;
+
+/**
+ * security data
+ */
+@Data
+public class SecurityData {
+
+    @NotNull
+    @Schema(description = "Conversation ID", example = "123")
+    private Long conversationId;
+
+    @NotNull
+    @Schema(description = "security data", example = "{\"password\":\"xxxxx\"}")
+    private String securityData;
+
+}

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

@@ -19,6 +19,7 @@
 package org.apache.hertzbeat.ai.service;
 
 import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
+import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
 import org.apache.hertzbeat.common.entity.ai.ChatConversation;
 import org.springframework.http.codec.ServerSentEvent;
 import reactor.core.publisher.Flux;
@@ -33,7 +34,7 @@ public interface ConversationService {
     /**
      * Send a message and receive a streaming response
      *
-     * @param message The user's message
+     * @param message        The user's message
      * @param conversationId Optional conversation ID for continuing a chat
      * @return Flux of ServerSentEvent for streaming the response
      */
@@ -67,4 +68,13 @@ public interface ConversationService {
      * @param conversationId Conversation ID to delete
      */
     void deleteConversation(Long conversationId);
+
+    /**
+     * save security data for a conversation
+     *
+     * @param securityData   securityData
+     * @return save result
+     */
+    Boolean saveSecurityData(SecurityData securityData);
+
 }

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

@@ -18,6 +18,10 @@
 
 package org.apache.hertzbeat.ai.service.impl;
 
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.hertzbeat.ai.sop.model.SopDefinition;
 import org.apache.hertzbeat.ai.sop.model.SopParameter;
@@ -27,9 +31,12 @@
 import org.apache.hertzbeat.ai.service.ChatClientProviderService;
 import org.apache.hertzbeat.base.dao.GeneralConfigDao;
 import org.apache.hertzbeat.common.entity.manager.GeneralConfig;
+import org.apache.hertzbeat.common.support.event.AiProviderConfigChangeEvent;
 import org.apache.hertzbeat.common.util.JsonUtil;
+import org.springframework.ai.chat.prompt.SystemPromptTemplate;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.context.event.EventListener;
 import org.springframework.core.io.Resource;
 import org.springframework.stereotype.Service;
 import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
@@ -44,14 +51,12 @@
 import reactor.core.publisher.Flux;
 
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 
 /**
- * Implementation of the {@link ChatClientProviderService}.
- * Provides functionality to interact with the ChatClient for handling chat
- * messages.
+ * Implementation of the {@link ChatClientProviderService}. Provides functionality to interact with the ChatClient for
+ * handling chat messages.
  */
 @Slf4j
 @Service
@@ -64,21 +69,28 @@ public class ChatClientProviderServiceImpl implements ChatClientProviderService
 
     private final GeneralConfigDao generalConfigDao;
 
+    private ModelProviderConfig modelProviderConfig;
+
+
     private final SkillRegistry skillRegistry;
-    
+
     @Autowired
     @Qualifier("hertzbeatTools")
     private ToolCallbackProvider toolCallbackProvider;
-    
+
     private boolean isConfigured = false;
 
     @Value("classpath:/prompt/system-message.st")
     private Resource systemResource;
 
+    @Value("classpath:/prompt/extra-message-protected.st")
+    private Resource extraResourceProtected;
+
+
     @Autowired
-    public ChatClientProviderServiceImpl(ApplicationContext applicationContext, 
-                                         GeneralConfigDao generalConfigDao,
-                                         @Lazy SkillRegistry skillRegistry) {
+    public ChatClientProviderServiceImpl(ApplicationContext applicationContext,
+        GeneralConfigDao generalConfigDao,
+        @Lazy SkillRegistry skillRegistry) {
         this.applicationContext = applicationContext;
         this.generalConfigDao = generalConfigDao;
         this.skillRegistry = skillRegistry;
@@ -89,7 +101,7 @@ public Flux<String> streamChat(ChatRequestContext context) {
         try {
             // Get the current (potentially refreshed) ChatClient instance
             ChatClient chatClient = applicationContext.getBean("openAiChatClient", ChatClient.class);
-            
+
             List<Message> messages = new ArrayList<>();
 
             // Add conversation history if available
@@ -112,13 +124,13 @@ public Flux<String> streamChat(ChatRequestContext context) {
             String systemPrompt = buildSystemPrompt(context.getConversationId());
 
             return chatClient.prompt()
-                    .messages(messages)
-                    .system(systemPrompt)
-                    .toolCallbacks(toolCallbackProvider)
-                    .stream()
-                    .content()
-                    .doOnComplete(() -> log.info("Streaming completed for conversation: {}", context.getConversationId()))
-                    .doOnError(error -> log.error("Error in streaming chat: {}", error.getMessage(), error));
+                .messages(messages)
+                .system(systemPrompt)
+                .toolCallbacks(toolCallbackProvider)
+                .stream()
+                .content()
+                .doOnComplete(() -> log.info("Streaming completed for conversation: {}", context.getConversationId()))
+                .doOnError(error -> log.error("Error in streaming chat: {}", error.getMessage(), error));
 
         } catch (Exception e) {
             log.error("Error setting up streaming chat: {}", e.getMessage(), e);
@@ -133,30 +145,43 @@ private String buildSystemPrompt(Long conversationId) {
         try {
             String template = systemResource.getContentAsString(StandardCharsets.UTF_8);
             String skillsList = generateSkillsList();
-            return template
-                    .replace(SKILLS_PLACEHOLDER, skillsList)
-                    .replace(CONVERSATION_ID_PLACEHOLDER, String.valueOf(conversationId));
+            template = template
+                .replace(SKILLS_PLACEHOLDER, skillsList)
+                .replace(CONVERSATION_ID_PLACEHOLDER, String.valueOf(conversationId));
+
+            // add extra prompt for protected model to guide it to use protected tools
+            if (Objects.equals(modelProviderConfig.getParticipationModel(), "PROTECTED")) {
+                Map<String, Object> metadata = new HashMap<>();
+                metadata.put("conversationId", conversationId);
+                return template + SystemPromptTemplate.builder().resource(extraResourceProtected).build()
+                    .create(metadata)
+                    .getContents();
+            } else {
+                return template;
+            }
+
         } catch (IOException e) {
             log.error("Failed to read system prompt template: {}", e.getMessage());
             return "";
         }
     }
 
+
     /**
      * Generate a formatted list of available skills for the system prompt.
      */
     private String generateSkillsList() {
         List<SopDefinition> skills = skillRegistry.getAllSkills();
-        
+
         if (skills.isEmpty()) {
             return "No skills currently available. Use listSkills tool to refresh.";
         }
-        
+
         StringBuilder sb = new StringBuilder();
         for (SopDefinition skill : skills) {
             sb.append("- **").append(skill.getName()).append("**: ");
             sb.append(skill.getDescription());
-            
+
             // Add parameter hints
             if (skill.getParameters() != null && !skill.getParameters().isEmpty()) {
                 sb.append(" (requires: ");
@@ -171,16 +196,24 @@ private String generateSkillsList() {
             }
             sb.append("\n");
         }
-        
+
         return sb.toString();
     }
 
+    @EventListener(AiProviderConfigChangeEvent.class)
+    public void onAiProviderConfigChange(AiProviderConfigChangeEvent event) {
+        GeneralConfig providerConfig = generalConfigDao.findByType("provider");
+        this.modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
+    }
+
     @Override
     public boolean isConfigured() {
         if (!isConfigured) {
             GeneralConfig providerConfig = generalConfigDao.findByType("provider");
-            ModelProviderConfig modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
-            isConfigured = modelProviderConfig != null && modelProviderConfig.getApiKey() != null;   
+            ModelProviderConfig modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(),
+                ModelProviderConfig.class);
+            isConfigured = modelProviderConfig != null && modelProviderConfig.getApiKey() != null;
+            this.modelProviderConfig = modelProviderConfig;
         }
         return isConfigured;
     }