From a2e7ff293bb5d86619f7aaa3438af9d5c538b48f Mon Sep 17 00:00:00 2001 From: Dmytro Skarzhynets Date: Wed, 2 Jul 2025 19:46:02 +0300 Subject: [PATCH] AI rule node: add API for sending arbitrary chat requests to LLM model --- .../server/actors/ActorSystemContext.java | 5 -- .../actors/ruleChain/DefaultTbContext.java | 6 -- .../server/controller/AiModelController.java | 63 +++++++++++++++ .../server/service/ai/AiModelService.java | 20 +++++ .../server/service/ai/AiModelServiceImpl.java | 12 ++- .../service/ai}/AiRequestsExecutor.java | 2 +- .../service/ai/DefaultAiRequestsExecutor.java | 1 - .../common/data/ai/dto/TbChatRequest.java | 79 +++++++++++++++++++ .../common/data/ai/dto/TbChatResponse.java | 68 ++++++++++++++++ .../server/common/data/ai/dto/TbContent.java | 73 +++++++++++++++++ .../common/data/ai/dto/TbUserMessage.java | 32 ++++++++ .../engine/api/RuleEngineAiModelService.java | 6 +- .../rule/engine/api/TbContext.java | 2 - .../thingsboard/rule/engine/ai/TbAiNode.java | 41 +++++----- 14 files changed, 367 insertions(+), 43 deletions(-) create mode 100644 application/src/main/java/org/thingsboard/server/controller/AiModelController.java create mode 100644 application/src/main/java/org/thingsboard/server/service/ai/AiModelService.java rename {rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api => application/src/main/java/org/thingsboard/server/service/ai}/AiRequestsExecutor.java (95%) create mode 100644 common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatRequest.java create mode 100644 common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatResponse.java create mode 100644 common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbContent.java create mode 100644 common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbUserMessage.java diff --git a/application/src/main/java/org/thingsboard/server/actors/ActorSystemContext.java b/application/src/main/java/org/thingsboard/server/actors/ActorSystemContext.java index 02b0164272..cd4a88314b 100644 --- a/application/src/main/java/org/thingsboard/server/actors/ActorSystemContext.java +++ b/application/src/main/java/org/thingsboard/server/actors/ActorSystemContext.java @@ -30,7 +30,6 @@ import org.springframework.context.annotation.Lazy; import org.springframework.data.redis.core.RedisTemplate; import org.springframework.stereotype.Component; import org.thingsboard.common.util.JacksonUtil; -import org.thingsboard.rule.engine.api.AiRequestsExecutor; import org.thingsboard.rule.engine.api.DeviceStateManager; import org.thingsboard.rule.engine.api.JobManager; import org.thingsboard.rule.engine.api.MailService; @@ -322,10 +321,6 @@ public class ActorSystemContext { @Getter private AiModelSettingsService aiModelSettingsService; - @Autowired - @Getter - private AiRequestsExecutor aiRequestsExecutor; - @Autowired @Getter private EntityViewService entityViewService; diff --git a/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java b/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java index d650b2dd21..b4235c7262 100644 --- a/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java +++ b/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java @@ -23,7 +23,6 @@ import org.bouncycastle.util.Arrays; import org.thingsboard.common.util.DebugModeUtil; import org.thingsboard.common.util.JacksonUtil; import org.thingsboard.common.util.ListeningExecutor; -import org.thingsboard.rule.engine.api.AiRequestsExecutor; import org.thingsboard.rule.engine.api.DeviceStateManager; import org.thingsboard.rule.engine.api.JobManager; import org.thingsboard.rule.engine.api.MailService; @@ -1037,11 +1036,6 @@ public class DefaultTbContext implements TbContext { return mainCtx.getAiModelSettingsService(); } - @Override - public AiRequestsExecutor getAiRequestsExecutor() { - return mainCtx.getAiRequestsExecutor(); - } - @Override public MqttClientSettings getMqttClientSettings() { return mainCtx.getMqttClientSettings(); diff --git a/application/src/main/java/org/thingsboard/server/controller/AiModelController.java b/application/src/main/java/org/thingsboard/server/controller/AiModelController.java new file mode 100644 index 0000000000..b9322052a3 --- /dev/null +++ b/application/src/main/java/org/thingsboard/server/controller/AiModelController.java @@ -0,0 +1,63 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.controller; + +import com.google.common.util.concurrent.ListenableFuture; +import dev.langchain4j.model.chat.request.ChatRequest; +import jakarta.validation.Valid; +import lombok.RequiredArgsConstructor; +import org.springframework.security.access.prepost.PreAuthorize; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.context.request.async.DeferredResult; +import org.thingsboard.server.common.data.ai.dto.TbChatRequest; +import org.thingsboard.server.common.data.ai.dto.TbChatResponse; +import org.thingsboard.server.common.data.ai.model.chat.AiChatModel; +import org.thingsboard.server.config.annotations.ApiOperation; +import org.thingsboard.server.service.ai.AiModelService; + +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static org.thingsboard.server.controller.ControllerConstants.TENANT_AUTHORITY_PARAGRAPH; + +@RestController +@RequiredArgsConstructor +@RequestMapping("/api/ai/model") +class AiModelController extends BaseController { + + private final AiModelService aiModelService; + + @ApiOperation( + value = "Send request to AI chat model (sendChatRequest)", + notes = "Submits a single prompt - made up of an optional system message and a required user message - to the specified AI chat model " + + "and returns either the generated answer or an error envelope." + + TENANT_AUTHORITY_PARAGRAPH + ) + @PreAuthorize("hasAuthority('TENANT_ADMIN')") + @PostMapping("/chat") + public DeferredResult sendChatRequest(@Valid @RequestBody TbChatRequest tbChatRequest) { + ChatRequest langChainChatRequest = tbChatRequest.toLangChainChatRequest(); + AiChatModel chatModel = tbChatRequest.chatModel(); + + ListenableFuture future = aiModelService.sendChatRequestAsync(chatModel, langChainChatRequest) + .transform(chatResponse -> (TbChatResponse) new TbChatResponse.Success(chatResponse.aiMessage().text()), directExecutor()) + .catching(Throwable.class, ex -> new TbChatResponse.Failure(ex.getMessage()), directExecutor()); + + return wrapFuture(future); + } + +} diff --git a/application/src/main/java/org/thingsboard/server/service/ai/AiModelService.java b/application/src/main/java/org/thingsboard/server/service/ai/AiModelService.java new file mode 100644 index 0000000000..196a53ce3a --- /dev/null +++ b/application/src/main/java/org/thingsboard/server/service/ai/AiModelService.java @@ -0,0 +1,20 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.service.ai; + +import org.thingsboard.rule.engine.api.RuleEngineAiModelService; + +public interface AiModelService extends RuleEngineAiModelService {} diff --git a/application/src/main/java/org/thingsboard/server/service/ai/AiModelServiceImpl.java b/application/src/main/java/org/thingsboard/server/service/ai/AiModelServiceImpl.java index 728e46e922..877f80fb26 100644 --- a/application/src/main/java/org/thingsboard/server/service/ai/AiModelServiceImpl.java +++ b/application/src/main/java/org/thingsboard/server/service/ai/AiModelServiceImpl.java @@ -15,23 +15,27 @@ */ package org.thingsboard.server.service.ai; +import com.google.common.util.concurrent.FluentFuture; import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; import lombok.RequiredArgsConstructor; import org.springframework.stereotype.Service; -import org.thingsboard.rule.engine.api.RuleEngineAiModelService; import org.thingsboard.server.common.data.ai.model.chat.AiChatModel; import org.thingsboard.server.common.data.ai.model.chat.AiChatModelConfig; import org.thingsboard.server.common.data.ai.model.chat.Langchain4jChatModelConfigurer; @Service @RequiredArgsConstructor -class AiModelServiceImpl implements RuleEngineAiModelService { +class AiModelServiceImpl implements AiModelService { private final Langchain4jChatModelConfigurer chatModelConfigurer; + private final AiRequestsExecutor aiRequestsExecutor; @Override - public > ChatModel configureChatModel(AiChatModel chatModel) { - return chatModel.configure(chatModelConfigurer); + public > FluentFuture sendChatRequestAsync(AiChatModel chatModel, ChatRequest chatRequest) { + ChatModel lc4jChatModel = chatModel.configure(chatModelConfigurer); + return aiRequestsExecutor.sendChatRequestAsync(lc4jChatModel, chatRequest); } } diff --git a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/AiRequestsExecutor.java b/application/src/main/java/org/thingsboard/server/service/ai/AiRequestsExecutor.java similarity index 95% rename from rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/AiRequestsExecutor.java rename to application/src/main/java/org/thingsboard/server/service/ai/AiRequestsExecutor.java index 9fd2a000ce..75de36c36d 100644 --- a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/AiRequestsExecutor.java +++ b/application/src/main/java/org/thingsboard/server/service/ai/AiRequestsExecutor.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.thingsboard.rule.engine.api; +package org.thingsboard.server.service.ai; import com.google.common.util.concurrent.FluentFuture; import dev.langchain4j.model.chat.ChatModel; diff --git a/application/src/main/java/org/thingsboard/server/service/ai/DefaultAiRequestsExecutor.java b/application/src/main/java/org/thingsboard/server/service/ai/DefaultAiRequestsExecutor.java index 077efa227d..2d0121d30a 100644 --- a/application/src/main/java/org/thingsboard/server/service/ai/DefaultAiRequestsExecutor.java +++ b/application/src/main/java/org/thingsboard/server/service/ai/DefaultAiRequestsExecutor.java @@ -33,7 +33,6 @@ import org.springframework.context.annotation.Lazy; import org.springframework.stereotype.Component; import org.springframework.validation.annotation.Validated; import org.thingsboard.common.util.ThingsBoardThreadFactory; -import org.thingsboard.rule.engine.api.AiRequestsExecutor; import java.time.Duration; import java.util.concurrent.Executors; diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatRequest.java b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatRequest.java new file mode 100644 index 0000000000..aa737b27d3 --- /dev/null +++ b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatRequest.java @@ -0,0 +1,79 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.common.data.ai.dto; + +import dev.langchain4j.data.message.ChatMessage; +import dev.langchain4j.data.message.Content; +import dev.langchain4j.data.message.SystemMessage; +import dev.langchain4j.data.message.UserMessage; +import dev.langchain4j.model.chat.request.ChatRequest; +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; +import org.thingsboard.server.common.data.ai.model.chat.AiChatModel; + +import java.util.ArrayList; +import java.util.List; + +public record TbChatRequest( + @Schema( + requiredMode = Schema.RequiredMode.NOT_REQUIRED, + accessMode = Schema.AccessMode.READ_WRITE, + description = "A system-level instruction that frames the user's input, setting the persona, tone, and constraints for the generated response", + example = "You are a helpful assistant. Only output valid JSON." + ) + String systemMessage, + + @Schema( + requiredMode = Schema.RequiredMode.REQUIRED, + accessMode = Schema.AccessMode.READ_WRITE, + description = "The actual user prompt that will be answered by the AI model" + ) + @NotNull @Valid + TbUserMessage userMessage, + + @Schema( + requiredMode = Schema.RequiredMode.REQUIRED, + accessMode = Schema.AccessMode.READ_WRITE, + description = "Configuration of the AI chat model that should execute the request" + ) + @NotNull @Valid + AiChatModel chatModel +) { + + public ChatRequest toLangChainChatRequest() { + return ChatRequest.builder() + .messages(getLangChainMessages()) + .build(); + } + + private List getLangChainMessages() { + List messages = new ArrayList<>(2); + + if (systemMessage != null) { + messages.add(SystemMessage.from(systemMessage)); + } + + List langChainContents = userMessage.contents().stream() + .map(TbContent::toLangChainContent) + .toList(); + + messages.add(UserMessage.from(langChainContents)); + + return messages; + } + +} diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatResponse.java b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatResponse.java new file mode 100644 index 0000000000..2cc17e4553 --- /dev/null +++ b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbChatResponse.java @@ -0,0 +1,68 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.common.data.ai.dto; + +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; +import io.swagger.v3.oas.annotations.media.Schema; + +@JsonTypeInfo( + use = JsonTypeInfo.Id.NAME, + property = "status", + include = JsonTypeInfo.As.PROPERTY, + visible = true +) +@JsonSubTypes({ + @JsonSubTypes.Type(value = TbChatResponse.Success.class, name = "SUCCESS"), + @JsonSubTypes.Type(value = TbChatResponse.Failure.class, name = "FAILURE") +}) +public sealed interface TbChatResponse permits TbChatResponse.Success, TbChatResponse.Failure { + + @Schema( + description = "Indicates whether the request was successful or not", + example = "SUCCESS" + ) + String getStatus(); + + record Success( + @Schema(description = "The text content generated by the model") + String generatedContent + ) implements TbChatResponse { + + @Override + @Schema(example = "SUCCESS") + public String getStatus() { + return "SUCCESS"; + } + + } + + record Failure( + @Schema( + description = "A string containing details about the failure" + ) + String errorDetails + ) implements TbChatResponse { + + @Override + @Schema(example = "FAILURE") + public String getStatus() { + return "FAILURE"; + } + + } + +} diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbContent.java b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbContent.java new file mode 100644 index 0000000000..23543121a4 --- /dev/null +++ b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbContent.java @@ -0,0 +1,73 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.common.data.ai.dto; + +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; +import dev.langchain4j.data.message.Content; +import dev.langchain4j.data.message.TextContent; +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.constraints.NotBlank; + +import static org.thingsboard.server.common.data.ai.dto.TbContent.TbTextContent; + +@JsonTypeInfo( + use = JsonTypeInfo.Id.NAME, + include = JsonTypeInfo.As.PROPERTY, + property = "contentType", + visible = true +) +@JsonSubTypes({ + @JsonSubTypes.Type(value = TbTextContent.class, name = "TEXT") +}) +public sealed interface TbContent permits TbTextContent { + + TbContentType contentType(); + + Content toLangChainContent(); + + enum TbContentType { + + TEXT + + } + + @Schema( + description = "Text-based content part of a user's prompt" + ) + record TbTextContent( + @NotBlank + @Schema( + requiredMode = Schema.RequiredMode.REQUIRED, + description = "The text content", + example = "What is the weather like in Kyiv today?" + ) + String text + ) implements TbContent { + + @Override + public TbContentType contentType() { + return TbContentType.TEXT; + } + + @Override + public Content toLangChainContent() { + return TextContent.from(text); + } + + } + +} diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbUserMessage.java b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbUserMessage.java new file mode 100644 index 0000000000..fdd7d8dc63 --- /dev/null +++ b/common/data/src/main/java/org/thingsboard/server/common/data/ai/dto/TbUserMessage.java @@ -0,0 +1,32 @@ +/** + * Copyright © 2016-2025 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.common.data.ai.dto; + +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotEmpty; + +import java.util.List; + +public record TbUserMessage( + @NotEmpty + @Valid + @Schema( + requiredMode = Schema.RequiredMode.REQUIRED, + description = "A list of content parts that make up the complete user prompt" + ) + List contents +) {} diff --git a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineAiModelService.java b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineAiModelService.java index 7ca01cab25..965a5b0f58 100644 --- a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineAiModelService.java +++ b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineAiModelService.java @@ -15,12 +15,14 @@ */ package org.thingsboard.rule.engine.api; -import dev.langchain4j.model.chat.ChatModel; +import com.google.common.util.concurrent.FluentFuture; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; import org.thingsboard.server.common.data.ai.model.chat.AiChatModel; import org.thingsboard.server.common.data.ai.model.chat.AiChatModelConfig; public interface RuleEngineAiModelService { - > ChatModel configureChatModel(AiChatModel chatModel); + > FluentFuture sendChatRequestAsync(AiChatModel chatModel, ChatRequest chatRequest); } diff --git a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java index a28695603f..1eeb644bba 100644 --- a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java +++ b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java @@ -427,8 +427,6 @@ public interface TbContext { AiModelSettingsService getAiModelSettingsService(); - AiRequestsExecutor getAiRequestsExecutor(); - // Configuration parameters for the MQTT client that is used in the MQTT node and Azure IoT hub node MqttClientSettings getMqttClientSettings(); diff --git a/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java b/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java index 75faf00c04..9e7d3fb23b 100644 --- a/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java +++ b/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java @@ -21,7 +21,6 @@ import com.google.common.util.concurrent.FutureCallback; import dev.langchain4j.data.message.ChatMessage; import dev.langchain4j.data.message.SystemMessage; import dev.langchain4j.data.message.UserMessage; -import dev.langchain4j.model.chat.ChatModel; import dev.langchain4j.model.chat.request.ChatRequest; import dev.langchain4j.model.chat.request.ResponseFormat; import dev.langchain4j.model.chat.request.ResponseFormatType; @@ -119,29 +118,27 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode { .responseFormat(responseFormat) .build(); - configureChatModelAsync(ctx) - .transformAsync(chatModel -> ctx.getAiRequestsExecutor().sendChatRequestAsync(chatModel, chatRequest), directExecutor()) - .addCallback(new FutureCallback<>() { - @Override - public void onSuccess(ChatResponse chatResponse) { - String response = chatResponse.aiMessage().text(); - if (!isValidJsonObject(response)) { - response = wrapInJsonObject(response); - } - tellSuccess(ctx, ackedMsg.transform() - .data(response) - .build()); - } + sendChatRequestAsync(ctx, chatRequest).addCallback(new FutureCallback<>() { + @Override + public void onSuccess(ChatResponse chatResponse) { + String response = chatResponse.aiMessage().text(); + if (!isValidJsonObject(response)) { + response = wrapInJsonObject(response); + } + tellSuccess(ctx, ackedMsg.transform() + .data(response) + .build()); + } - @Override - public void onFailure(@NonNull Throwable t) { - tellFailure(ctx, ackedMsg, t); - } - }, directExecutor()); + @Override + public void onFailure(@NonNull Throwable t) { + tellFailure(ctx, ackedMsg, t); + } + }, directExecutor()); } - private > FluentFuture configureChatModelAsync(TbContext ctx) { - return ctx.getAiModelSettingsService().findAiModelSettingsByTenantIdAndIdAsync(ctx.getTenantId(), modelSettingsId).transform(settingsOpt -> { + private > FluentFuture sendChatRequestAsync(TbContext ctx, ChatRequest chatRequest) { + return ctx.getAiModelSettingsService().findAiModelSettingsByTenantIdAndIdAsync(ctx.getTenantId(), modelSettingsId).transformAsync(settingsOpt -> { if (settingsOpt.isEmpty()) { throw new NoSuchElementException("[" + ctx.getTenantId() + "] AI model settings with ID: [" + modelSettingsId + "] were not found"); } @@ -158,7 +155,7 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode { .withTimeoutSeconds(timeoutSeconds) .withMaxRetries(0)); // disable retries to respect timeout set in rule node config - return ctx.getAiModelService().configureChatModel(chatModel); + return ctx.getAiModelService().sendChatRequestAsync(chatModel, chatRequest); }, ctx.getDbCallbackExecutor()); }