From 19c234fcdf58175cd1d9401a91c194ac4c01081c Mon Sep 17 00:00:00 2001
From: Dmytro Skarzhynets <dskarzhynets@thingsboard.io>
Date: Mon, 9 Jun 2025 17:51:24 +0300
Subject: [PATCH] AI rule node: configure chat model on each message

---
 .../org/thingsboard/rule/engine/ai/TbAiNode.java    | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

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 564de8d66e..6b0bb3803b 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
@@ -35,6 +35,7 @@ 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.id.AiSettingsId;
 import org.thingsboard.server.common.data.plugin.ComponentType;
 import org.thingsboard.server.common.msg.TbMsg;
 import org.thingsboard.server.dao.exception.DataValidationException;
@@ -59,7 +60,7 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
     private SystemMessage systemMessage;
     private PromptTemplate userPromptTemplate;
     private ResponseFormat responseFormat;
-    private ChatModel chatModel;
+    private AiSettingsId aiSettingsId;
 
     @Override
     public void init(TbContext ctx, TbNodeConfiguration configuration) throws TbNodeException {
@@ -86,7 +87,8 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
                 Rule engine message type: {{msgType}}"""
                 .formatted(config.getUserPrompt())
         );
-        chatModel = ctx.getAiService().configureChatModel(ctx.getTenantId(), config.getAiSettingsId());
+
+        aiSettingsId = config.getAiSettingsId();
     }
 
     private static JsonSchema getJsonSchema(ResponseFormatType responseFormatType, JsonNode jsonSchema) {
@@ -112,7 +114,9 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
                 .responseFormat(responseFormat)
                 .build();
 
-        addCallback(sendChatRequest(ctx, chatRequest), new FutureCallback<>() {
+        ChatModel chatModel = ctx.getAiService().configureChatModel(ctx.getTenantId(), aiSettingsId);
+
+        addCallback(sendChatRequest(ctx, chatModel, chatRequest), new FutureCallback<>() {
             @Override
             public void onSuccess(String response) {
                 if (!isValidJsonObject(response)) {
@@ -130,7 +134,7 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
         }, directExecutor());
     }
 
-    private ListenableFuture<String> sendChatRequest(TbContext ctx, ChatRequest chatRequest) {
+    private ListenableFuture<String> sendChatRequest(TbContext ctx, ChatModel chatModel, ChatRequest chatRequest) {
         return ctx.getExternalCallExecutor().submit(() -> chatModel.chat(chatRequest).aiMessage().text());
     }
 
@@ -153,7 +157,6 @@ public final class TbAiNode extends TbAbstractExternalNode implements TbNode {
         systemMessage = null;
         userPromptTemplate = null;
         responseFormat = null;
-        chatModel = null;
     }
 
 }