tsai
/
thingsboard
mirror of https://github.com/thingsboard/thingsboard.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

AI rule node: make message processing non-blocking

pull/13371/head
Dmytro Skarzhynets 1 year ago
parent
fe9caaedd4
commit
db5e4f8d91
No known key found for this signature in database GPG Key ID: 2B51652F224037DF
3 changed files with 44 additions and 34 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 3
      common/dao-api/src/main/java/org/thingsboard/server/dao/ai/AiSettingsService.java
  2. 8
      dao/src/main/java/org/thingsboard/server/dao/ai/AiSettingsServiceImpl.java
  3. 67
      rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java

3
common/dao-api/src/main/java/org/thingsboard/server/dao/ai/AiSettingsService.java
View File

@ -15,6 +15,7 @@
*/
package org.thingsboard.server.dao.ai;
import com.google.common.util.concurrent.FluentFuture;
import org.thingsboard.server.common.data.ai.AiSettings;
import org.thingsboard.server.common.data.id.AiSettingsId;
import org.thingsboard.server.common.data.id.TenantId;
@ -34,6 +35,8 @@ public interface AiSettingsService extends EntityDaoService {
Optional<AiSettings> findAiSettingsByTenantIdAndId(TenantId tenantId, AiSettingsId aiSettingsId);
FluentFuture<Optional<AiSettings>> findAiSettingsByTenantIdAndIdAsync(TenantId tenantId, AiSettingsId aiSettingsId);
boolean deleteByTenantIdAndId(TenantId tenantId, AiSettingsId aiSettingsId);
}

8
dao/src/main/java/org/thingsboard/server/dao/ai/AiSettingsServiceImpl.java
View File

@ -16,6 +16,7 @@
package org.thingsboard.server.dao.ai;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.FluentFuture;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
@ -32,6 +33,7 @@ import org.thingsboard.server.dao.entity.CachedVersionedEntityService;
import org.thingsboard.server.dao.eventsourcing.DeleteEntityEvent;
import org.thingsboard.server.dao.eventsourcing.SaveEntityEvent;
import org.thingsboard.server.dao.service.DataValidator;
import org.thingsboard.server.dao.sql.JpaExecutorService;
import java.util.List;
import java.util.Optional;
@ -45,6 +47,7 @@ class AiSettingsServiceImpl extends CachedVersionedEntityService<AiSettingsCache
private final DataValidator<AiSettings> aiSettingsValidator;
private final JpaExecutorService jpaExecutor;
private final AiSettingsDao aiSettingsDao;
@Override
@ -97,6 +100,11 @@ class AiSettingsServiceImpl extends CachedVersionedEntityService<AiSettingsCache
return Optional.ofNullable(cache.get(cacheKey, () -> aiSettingsDao.findByTenantIdAndId(tenantId, aiSettingsId).orElse(null)));
}
@Override
public FluentFuture<Optional<AiSettings>> findAiSettingsByTenantIdAndIdAsync(TenantId tenantId, AiSettingsId aiSettingsId) {
return FluentFuture.from(jpaExecutor.submit(() -> findAiSettingsByTenantIdAndId(tenantId, aiSettingsId)));
}
@Override
@Transactional
public boolean deleteByTenantIdAndId(TenantId tenantId, AiSettingsId aiSettingsId) {

67
rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/ai/TbAiNode.java
View File

@ -17,8 +17,8 @@ package org.thingsboard.rule.engine.ai;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.util.concurrent.FluentFuture;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.ListenableFuture;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
@ -36,7 +36,6 @@ import org.thingsboard.rule.engine.api.TbNodeConfiguration;
import org.thingsboard.rule.engine.api.TbNodeException;
import org.thingsboard.rule.engine.api.util.TbNodeUtils;
import org.thingsboard.rule.engine.external.TbAbstractExternalNode;
import org.thingsboard.server.common.data.ai.AiSettings;
import org.thingsboard.server.common.data.ai.model.AiModelConfig;
import org.thingsboard.server.common.data.ai.provider.AiProviderConfig;
import org.thingsboard.server.common.data.id.AiSettingsId;
@ -47,9 +46,8 @@ import org.thingsboard.server.dao.exception.DataValidationException;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.NoSuchElementException;
import static com.google.common.util.concurrent.Futures.addCallback;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.Objects.requireNonNullElse;
import static org.thingsboard.server.dao.service.ConstraintValidator.validateFields;
@ -131,43 +129,44 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
.responseFormat(responseFormat)
.build();
ChatModel chatModel = configureChatModel(ctx);
addCallback(sendChatRequest(ctx, chatModel, chatRequest), new FutureCallback<>() {
@Override
public void onSuccess(String response) {
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());
configureChatModelAsync(ctx)
.transform(chatModel -> sendChatRequest(chatModel, chatRequest), ctx.getExternalCallExecutor())
.addCallback(new FutureCallback<>() {
@Override
public void onSuccess(String response) {
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());
}
private ChatModel configureChatModel(TbContext ctx) throws TbNodeException {
Optional<AiSettings> aiSettingsOpt = ctx.getAiSettingsService().findAiSettingsByTenantIdAndId(ctx.getTenantId(), aiSettingsId);
if (aiSettingsOpt.isEmpty()) {
throw new TbNodeException("AI settings with ID: " + aiSettingsId + " were not found", true);
}
private FluentFuture<ChatModel> configureChatModelAsync(TbContext ctx) {
return ctx.getAiSettingsService().findAiSettingsByTenantIdAndIdAsync(ctx.getTenantId(), aiSettingsId).transform(aiSettingsOpt -> {
if (aiSettingsOpt.isEmpty()) {
throw new NoSuchElementException("AI settings with ID: " + aiSettingsId + " were not found");
}
AiProviderConfig providerConfig = aiSettingsOpt.get().getProviderConfig();
AiModelConfig modelConfig = aiSettingsOpt.get().getModelConfig();
AiProviderConfig providerConfig = aiSettingsOpt.get().getProviderConfig();
AiModelConfig modelConfig = aiSettingsOpt.get().getModelConfig();
modelConfig.setTimeoutSeconds(timeoutSeconds);
modelConfig.setMaxRetries(0); // disable retries to respect timeout set in rule node config
modelConfig.setTimeoutSeconds(timeoutSeconds);
modelConfig.setMaxRetries(0); // disable retries to respect timeout set in rule node config
return ctx.getAiService().configureChatModel(providerConfig, modelConfig);
return ctx.getAiService().configureChatModel(providerConfig, modelConfig);
}, ctx.getDbCallbackExecutor());
}
private ListenableFuture<String> sendChatRequest(TbContext ctx, ChatModel chatModel, ChatRequest chatRequest) {
return ctx.getExternalCallExecutor().submit(() -> chatModel.chat(chatRequest).aiMessage().text());
private String sendChatRequest(ChatModel chatModel, ChatRequest chatRequest) {
return chatModel.chat(chatRequest).aiMessage().text();
}
private static boolean isValidJsonObject(String jsonString) {

Write Preview
Loading…
Cancel
Save