Skip to content

Commit 9b31bae

Browse files
yuluo-yxtomsun28
yuluo-yx
and
tomsun28
authored
fix(ai): improve SOP execution and provider config handling (#4224)
Signed-off-by: yuluo-yx <yuluo08290126@gmail.com> Co-authored-by: Tomsun28 <tomsun28@outlook.com>
1 parent 0a38a73 commit 9b31bae

9 files changed

Lines changed: 544 additions & 93 deletions

File tree

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

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import org.springframework.beans.factory.annotation.Autowired;
4949
import org.springframework.beans.factory.annotation.Qualifier;
5050
import org.springframework.context.ApplicationContext;
51+
import org.springframework.util.StringUtils;
5152
import reactor.core.publisher.Flux;
5253

5354
import java.io.IOException;
@@ -69,7 +70,7 @@ public class ChatClientProviderServiceImpl implements ChatClientProviderService
6970

7071
private final GeneralConfigDao generalConfigDao;
7172

72-
private ModelProviderConfig modelProviderConfig;
73+
private volatile ModelProviderConfig modelProviderConfig;
7374

7475

7576
private final SkillRegistry skillRegistry;
@@ -78,7 +79,9 @@ public class ChatClientProviderServiceImpl implements ChatClientProviderService
7879
@Qualifier("hertzbeatTools")
7980
private ToolCallbackProvider toolCallbackProvider;
8081

81-
private boolean isConfigured = false;
82+
private volatile boolean configured;
83+
84+
private volatile boolean configurationLoaded;
8285

8386
@Value("classpath:/prompt/system-message.st")
8487
private Resource systemResource;
@@ -150,7 +153,8 @@ private String buildSystemPrompt(Long conversationId) {
150153
.replace(CONVERSATION_ID_PLACEHOLDER, String.valueOf(conversationId));
151154

152155
// add extra prompt for protected model to guide it to use protected tools
153-
if (Objects.equals(modelProviderConfig.getParticipationModel(), "PROTECTED")) {
156+
ModelProviderConfig currentConfig = modelProviderConfig;
157+
if (currentConfig != null && Objects.equals(currentConfig.getParticipationModel(), "PROTECTED")) {
154158
Map<String, Object> metadata = new HashMap<>();
155159
metadata.put("conversationId", conversationId);
156160
return template + SystemPromptTemplate.builder().resource(extraResourceProtected).build()
@@ -202,19 +206,32 @@ private String generateSkillsList() {
202206

203207
@EventListener(AiProviderConfigChangeEvent.class)
204208
public void onAiProviderConfigChange(AiProviderConfigChangeEvent event) {
205-
GeneralConfig providerConfig = generalConfigDao.findByType("provider");
206-
this.modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
209+
refreshProviderConfiguration();
207210
}
208211

209212
@Override
210213
public boolean isConfigured() {
211-
if (!isConfigured) {
212-
GeneralConfig providerConfig = generalConfigDao.findByType("provider");
213-
ModelProviderConfig modelProviderConfig = JsonUtil.fromJson(providerConfig.getContent(),
214-
ModelProviderConfig.class);
215-
isConfigured = modelProviderConfig != null && modelProviderConfig.getApiKey() != null;
216-
this.modelProviderConfig = modelProviderConfig;
214+
if (!configurationLoaded) {
215+
synchronized (this) {
216+
if (!configurationLoaded) {
217+
refreshProviderConfiguration();
218+
}
219+
}
220+
}
221+
return configured;
222+
}
223+
224+
/**
225+
* Atomically refreshes the configuration snapshot after enabling, disabling, or switching providers.
226+
*/
227+
private synchronized void refreshProviderConfiguration() {
228+
GeneralConfig providerConfig = generalConfigDao.findByType("provider");
229+
ModelProviderConfig refreshedConfig = null;
230+
if (providerConfig != null && StringUtils.hasText(providerConfig.getContent())) {
231+
refreshedConfig = JsonUtil.fromJson(providerConfig.getContent(), ModelProviderConfig.class);
217232
}
218-
return isConfigured;
233+
modelProviderConfig = refreshedConfig;
234+
configured = refreshedConfig != null && StringUtils.hasText(refreshedConfig.getApiKey());
235+
configurationLoaded = true;
219236
}
220237
}

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/sop/engine/SopEngineImpl.java‎

Lines changed: 60 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,6 @@
4141
@Service
4242
public class SopEngineImpl implements SopEngine {
4343

44-
// Thread-local context for each execution
45-
private static final ThreadLocal<Map<String, Object>> CONTEXT_BUS = ThreadLocal.withInitial(HashMap::new);
46-
4744
private final List<SopExecutor> executors;
4845

4946
@Autowired
@@ -55,22 +52,12 @@ public SopEngineImpl(List<SopExecutor> executors) {
5552
public Flux<String> execute(SopDefinition definition, Map<String, Object> inputParams) {
5653
return Flux.create(sink -> {
5754
try {
55+
OutputConfig outputConfig = resolveOutputConfig(definition);
56+
Map<String, Object> context = prepareContext(definition, inputParams, outputConfig);
57+
5858
log.info("Starting execution of SOP: {}", definition.getName());
5959
sink.next("Starting SOP: " + definition.getName() + " (v" + definition.getVersion() + ")");
6060

61-
// Initialize context with input parameters
62-
Map<String, Object> context = CONTEXT_BUS.get();
63-
context.clear();
64-
context.putAll(inputParams);
65-
66-
// Add language configuration to context
67-
OutputConfig outputConfig = definition.getOutput();
68-
if (outputConfig != null) {
69-
context.put("_language", outputConfig.getLanguageCode());
70-
} else {
71-
context.put("_language", "zh");
72-
}
73-
7461
for (SopStep step : definition.getSteps()) {
7562
sink.next("Executing step [" + step.getId() + "]: " + step.getType());
7663

@@ -106,10 +93,9 @@ public Flux<String> execute(SopDefinition definition, Map<String, Object> inputP
10693
sink.next("SOP " + definition.getName() + " completed successfully.");
10794
sink.complete();
10895
} catch (Exception e) {
109-
log.error("Error executing SOP {}: {}", definition.getName(), e.getMessage(), e);
96+
String sopName = definition == null ? "unknown" : definition.getName();
97+
log.error("Error executing SOP {}: {}", sopName, e.getMessage(), e);
11098
sink.error(e);
111-
} finally {
112-
CONTEXT_BUS.remove();
11399
}
114100
});
115101
}
@@ -118,40 +104,20 @@ public Flux<String> execute(SopDefinition definition, Map<String, Object> inputP
118104
public SopResult executeSync(SopDefinition definition, Map<String, Object> inputParams) {
119105
long startTime = System.currentTimeMillis();
120106
List<StepResult> stepResults = new ArrayList<>();
121-
Map<String, Object> context = new HashMap<>(inputParams);
122-
123-
// Apply default values for parameters that are not provided
124-
if (definition.getParameters() != null) {
125-
for (SopParameter param : definition.getParameters()) {
126-
if (!context.containsKey(param.getName()) && param.getDefaultValue() != null) {
127-
context.put(param.getName(), param.getDefaultValue());
128-
}
129-
}
130-
}
131-
132-
// Get output configuration first
133-
OutputConfig outputConfig = definition.getOutput();
134-
if (outputConfig == null) {
135-
outputConfig = OutputConfig.builder()
136-
.type("simple")
137-
.format("text")
138-
.language("zh")
139-
.build();
140-
}
141-
142-
// Add language configuration to context
143-
context.put("_language", outputConfig.getLanguageCode());
144-
107+
145108
SopResult.SopResultBuilder resultBuilder = SopResult.builder()
146-
.sopName(definition.getName())
147-
.sopVersion(definition.getVersion())
109+
.sopName(definition == null ? null : definition.getName())
110+
.sopVersion(definition == null ? null : definition.getVersion())
148111
.startTime(startTime);
149-
150-
resultBuilder.outputType(outputConfig.getOutputType());
151-
resultBuilder.outputFormat(outputConfig.getFormat() != null ? outputConfig.getFormat() : "text");
152-
resultBuilder.language(outputConfig.getLanguageCode());
153-
112+
154113
try {
114+
OutputConfig outputConfig = resolveOutputConfig(definition);
115+
Map<String, Object> context = prepareContext(definition, inputParams, outputConfig);
116+
117+
resultBuilder.outputType(outputConfig.getOutputType());
118+
resultBuilder.outputFormat(outputConfig.getFormat() != null ? outputConfig.getFormat() : "text");
119+
resultBuilder.language(outputConfig.getLanguageCode());
120+
155121
log.info("Starting sync execution of SOP: {}", definition.getName());
156122

157123
for (SopStep step : definition.getSteps()) {
@@ -217,6 +183,50 @@ public SopResult executeSync(SopDefinition definition, Map<String, Object> input
217183
return buildFailedResult(resultBuilder, stepResults, e.getMessage(), startTime);
218184
}
219185
}
186+
187+
/**
188+
* Applies defaults and validates required parameters consistently for every execution entry point.
189+
*/
190+
private Map<String, Object> prepareContext(SopDefinition definition, Map<String, Object> inputParams,
191+
OutputConfig outputConfig) {
192+
if (definition.getSteps() == null || definition.getSteps().isEmpty()) {
193+
throw new IllegalArgumentException("SOP must contain at least one step");
194+
}
195+
196+
Map<String, Object> context = inputParams == null ? new HashMap<>() : new HashMap<>(inputParams);
197+
if (definition.getParameters() != null) {
198+
for (SopParameter parameter : definition.getParameters()) {
199+
Object value = context.get(parameter.getName());
200+
if (isMissing(value) && parameter.getDefaultValue() != null) {
201+
context.put(parameter.getName(), parameter.getDefaultValue());
202+
value = parameter.getDefaultValue();
203+
}
204+
if (parameter.isRequired() && isMissing(value)) {
205+
throw new IllegalArgumentException("Required SOP parameter is missing: " + parameter.getName());
206+
}
207+
}
208+
}
209+
context.put("_language", outputConfig.getLanguageCode());
210+
return context;
211+
}
212+
213+
private boolean isMissing(Object value) {
214+
return value == null || value instanceof String text && text.isBlank();
215+
}
216+
217+
private OutputConfig resolveOutputConfig(SopDefinition definition) {
218+
if (definition == null) {
219+
throw new IllegalArgumentException("SOP definition must not be null");
220+
}
221+
if (definition.getOutput() != null) {
222+
return definition.getOutput();
223+
}
224+
return OutputConfig.builder()
225+
.type("simple")
226+
.format("text")
227+
.language("zh")
228+
.build();
229+
}
220230

221231
private SopResult buildFailedResult(SopResult.SopResultBuilder builder,
222232
List<StepResult> stepResults, String error, long startTime) {

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/sop/registry/SopToolCallback.java‎

Lines changed: 61 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,16 @@
1717

1818
package org.apache.hertzbeat.ai.sop.registry;
1919

20+
import java.util.ArrayList;
21+
import java.util.HashMap;
22+
import java.util.LinkedHashMap;
23+
import java.util.List;
24+
import java.util.Map;
2025
import org.apache.hertzbeat.ai.sop.engine.SopEngine;
2126
import org.apache.hertzbeat.ai.sop.model.SopDefinition;
27+
import org.apache.hertzbeat.ai.sop.model.SopParameter;
28+
import org.apache.hertzbeat.ai.sop.model.SopResult;
29+
import org.apache.hertzbeat.common.util.JsonUtil;
2230
import org.springframework.ai.tool.ToolCallback;
2331
import org.springframework.ai.tool.definition.ToolDefinition;
2432

@@ -49,33 +57,66 @@ public ToolDefinition getToolDefinition() {
4957

5058
@Override
5159
public String call(String arguments) {
52-
// TODO: Parse arguments and execute SOP
53-
return "SOP " + definition.getName() + " execution started.";
60+
Map<String, Object> inputParams = parseArguments(arguments);
61+
SopResult result = sopEngine.executeSync(definition, inputParams);
62+
return result.toAiResponse();
5463
}
5564

5665
private String buildInputSchema() {
57-
// Build JSON Schema for SOP parameters
58-
StringBuilder schema = new StringBuilder();
59-
schema.append("{\"type\":\"object\",\"properties\":{");
60-
61-
if (definition.getParameters() != null && !definition.getParameters().isEmpty()) {
62-
boolean first = true;
63-
for (var param : definition.getParameters()) {
64-
if (!first) {
65-
schema.append(",");
66+
Map<String, Object> schema = new LinkedHashMap<>();
67+
Map<String, Object> properties = new LinkedHashMap<>();
68+
List<String> required = new ArrayList<>();
69+
70+
if (definition.getParameters() != null) {
71+
for (SopParameter parameter : definition.getParameters()) {
72+
Map<String, Object> property = new LinkedHashMap<>();
73+
property.put("type", mapType(parameter.getType()));
74+
if (parameter.getDescription() != null) {
75+
property.put("description", parameter.getDescription());
76+
}
77+
if (parameter.getDefaultValue() != null) {
78+
property.put("default", convertDefaultValue(parameter));
6679
}
67-
schema.append("\"").append(param.getName()).append("\":");
68-
schema.append("{\"type\":\"").append(mapType(param.getType())).append("\"");
69-
if (param.getDescription() != null) {
70-
schema.append(",\"description\":\"").append(param.getDescription()).append("\"");
80+
properties.put(parameter.getName(), property);
81+
if (parameter.isRequired()) {
82+
required.add(parameter.getName());
7183
}
72-
schema.append("}");
73-
first = false;
7484
}
7585
}
76-
77-
schema.append("}}");
78-
return schema.toString();
86+
87+
schema.put("type", "object");
88+
schema.put("properties", properties);
89+
if (!required.isEmpty()) {
90+
schema.put("required", required);
91+
}
92+
schema.put("additionalProperties", false);
93+
return JsonUtil.toJson(schema);
94+
}
95+
96+
@SuppressWarnings("unchecked")
97+
private Map<String, Object> parseArguments(String arguments) {
98+
if (arguments == null || arguments.isBlank()) {
99+
return new HashMap<>();
100+
}
101+
Map<String, Object> inputParams = JsonUtil.fromJson(arguments, Map.class);
102+
if (inputParams == null) {
103+
throw new IllegalArgumentException("SOP tool arguments must be a valid JSON object");
104+
}
105+
return inputParams;
106+
}
107+
108+
private Object convertDefaultValue(SopParameter parameter) {
109+
String defaultValue = parameter.getDefaultValue();
110+
try {
111+
return switch (mapType(parameter.getType())) {
112+
case "boolean" -> Boolean.valueOf(defaultValue);
113+
case "integer" -> Long.valueOf(defaultValue);
114+
case "number" -> Double.valueOf(defaultValue);
115+
default -> defaultValue;
116+
};
117+
} catch (NumberFormatException e) {
118+
throw new IllegalArgumentException("Invalid default value for SOP parameter: " + parameter.getName(), e);
119+
}
79120
}
80121

81122
private String mapType(String type) {

‎hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/sop/registry/ToolRegistry.java‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ public boolean hasMethod(String toolName) {
129129
*/
130130
public Set<String> getToolNames() {
131131
ensureInitialized();
132-
return tools.keySet();
132+
return Set.copyOf(tools.keySet());
133133
}
134134

135135
/**
@@ -176,15 +176,19 @@ private List<ParamInfo> extractParamInfos(Method method) {
176176
* Invoke this tool method with the given arguments.
177177
*/
178178
public String invoke(Map<String, Object> args) {
179-
try {
180-
Object[] methodArgs = new Object[paramInfos.size()];
181-
182-
for (int i = 0; i < paramInfos.size(); i++) {
183-
ParamInfo paramInfo = paramInfos.get(i);
184-
Object value = args.get(paramInfo.name);
185-
methodArgs[i] = convertValue(value, paramInfo.type);
179+
Map<String, Object> safeArgs = args == null ? Map.of() : args;
180+
Object[] methodArgs = new Object[paramInfos.size()];
181+
182+
for (int i = 0; i < paramInfos.size(); i++) {
183+
ParamInfo paramInfo = paramInfos.get(i);
184+
Object value = safeArgs.get(paramInfo.name);
185+
if (paramInfo.required && isMissing(value)) {
186+
throw new IllegalArgumentException("Required tool parameter is missing: " + paramInfo.name);
186187
}
187-
188+
methodArgs[i] = convertValue(value, paramInfo.type);
189+
}
190+
191+
try {
188192
Object result = method.invoke(bean, methodArgs);
189193
return result != null ? result.toString() : "";
190194

@@ -194,6 +198,10 @@ public String invoke(Map<String, Object> args) {
194198
}
195199
}
196200

201+
private boolean isMissing(Object value) {
202+
return value == null || value instanceof String text && text.isBlank();
203+
}
204+
197205
private Object convertValue(Object value, Class<?> targetType) {
198206
if (value == null) {
199207
return null;

0 commit comments

Comments
 (0)