QQL GenAi improvements

build.gradle

@@ -194,7 +194,10 @@ configure(javaProjects) {
 //                entry 'jackson-module-parameter-names'
 //            }
 
-            dependency 'com.azure:azure-security-keyvault-secrets:4.8.4'
+            dependency 'com.azure:azure-core:1.55.3'
+            dependency 'com.azure:azure-core-http-netty:1.15.11'
+            dependency 'com.azure:azure-identity:1.15.4'
+            dependency 'com.azure:azure-security-keyvault-secrets:4.9.4'
             dependency 'net.java.dev.jna:jna-platform:5.6.0' // azure sdk requires jna 5.6.0
             dependency 'com.microsoft.azure:msal4j:1.15.1'
 
@@ -284,6 +287,11 @@ configure(javaProjects) {
 
             dependency 'org.apache.bcel:bcel:6.6.0'
 
+            // AI integration
+            dependency 'dev.langchain4j:langchain4j:1.7.1'
+            dependency 'dev.langchain4j:langchain4j-azure-open-ai-spring-boot-starter:1.7.1-beta14'
+            dependency 'dev.langchain4j:langchain4j-spring-boot-starter:1.7.1-beta14'
+
             dependency 'org.reactivestreams:reactive-streams:1.0.4'
 
             dependencySet(group: 'io.netty', version: "$nettyVersion") {

java/ws-server/build.gradle

@@ -80,9 +80,10 @@ dependencies {
 
 //    implementation 'deltix:deltix-spring-api-keys'
 
-//    implementation 'io.reactivex.rxjava2:rxjava'
-//    implementation 'deltix:deltix-timebase-api-rx'
-//    implementation        'deltix:transformation-api'
+    // AI integration
+    implementation 'dev.langchain4j:langchain4j'
+    implementation 'dev.langchain4j:langchain4j-azure-open-ai-spring-boot-starter'
+    implementation 'dev.langchain4j:langchain4j-spring-boot-starter'
 
     compileOnly 'com.webcohesion.enunciate:enunciate-core-annotations'
 

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/config/GenAiConfig.java

@@ -0,0 +1,146 @@
+/*
+ * Copyright 2025 EPAM Systems, Inc
+ *
+ * See the NOTICE file distributed with this work for additional information
+ * regarding copyright ownership. 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 com.epam.deltix.tbwg.webapp.config;
+
+import com.epam.deltix.tbwg.webapp.services.genai.GenAiHelperService;
+import com.epam.deltix.tbwg.webapp.services.genai.models.*;
+import com.epam.deltix.tbwg.webapp.services.genai.docsloader.MarkdownCodeAwareSplitter;
+import com.epam.deltix.tbwg.webapp.services.genai.docsloader.MarkdownDocsLoader;
+import com.epam.deltix.tbwg.webapp.settings.AiApiSettings;
+import dev.langchain4j.data.document.Document;
+import dev.langchain4j.data.document.DocumentSplitter;
+import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.model.chat.ChatModel;
+import dev.langchain4j.model.chat.StreamingChatModel;
+import dev.langchain4j.model.embedding.EmbeddingModel;
+import dev.langchain4j.rag.content.retriever.ContentRetriever;
+import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
+import dev.langchain4j.service.AiServices;
+import dev.langchain4j.store.embedding.EmbeddingStore;
+import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
+import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.io.Resource;
+import org.springframework.core.io.ResourceLoader;
+import org.springframework.core.io.support.ResourcePatternResolver;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.StandardCopyOption;
+import java.util.List;
+
+@Configuration
+@ConditionalOnBean(AiApiSettings.class)
+public class GenAiConfig {
+
+    private static final String CLASSPATH_EMB_PATH = "classpath:qql_gen/qql-docs-embeddings.json";
+    private static final String FALLBACK_EMB_FILE = "qql-docs-embeddings.json";
+
+    @Bean
+    public PerUserAzureChatModel perUserAzureChatModel(AiApiSettings settings,
+                                                       UserAiApiKeyProvider keyProvider) {
+        return new PerUserAzureChatModel(settings.getEndpointUrl(),
+                settings.getDeploymentName(),
+                keyProvider);
+    }
+
+    @Bean
+    public PerUserAzureStreamingChatModel perUserAzureStreamingChatModel(AiApiSettings settings,
+                                                                         UserAiApiKeyProvider keyProvider) {
+        return new PerUserAzureStreamingChatModel(settings.getEndpointUrl(),
+                settings.getDeploymentName(),
+                keyProvider);
+    }
+
+    @Bean
+    public PerUserAzureEmbeddingModel perUserAzureEmbeddingModel(AiApiSettings settings,
+                                                                 UserAiApiKeyProvider keyProvider) {
+        return new PerUserAzureEmbeddingModel(settings.getEndpointUrl(),
+                settings.getEmbeddingDeploymentName(),
+                keyProvider);
+    }
+
+    @Bean
+    public GenAiHelperService genAiHelperService(ChatModel model, StreamingChatModel streamingModel) {
+        return AiServices.builder(GenAiHelperService.class)
+                .chatModel(model)
+                .streamingChatModel(streamingModel)
+                .build();
+    }
+
+    @Bean
+    public MarkdownDocsLoader markdownDocsLoader(ResourcePatternResolver resolver) {
+        return new MarkdownDocsLoader(resolver);
+    }
+
+    @Bean
+    public List<Document> qqlAllSectionDocs(MarkdownDocsLoader loader) throws IOException {
+        return loader.load();
+    }
+
+    @Bean
+    EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel,
+                                               MarkdownDocsLoader loader,
+                                               ResourceLoader resourceLoader) throws IOException {
+        Resource res = resourceLoader.getResource(CLASSPATH_EMB_PATH);
+        if (res.exists()) {
+            try (InputStream in = res.getInputStream()) {
+                var tmp = Files.createTempFile("qql-docs-embeddings", ".json");
+                Files.copy(in, tmp, StandardCopyOption.REPLACE_EXISTING);
+                return InMemoryEmbeddingStore.fromFile(tmp.toString());
+            }
+        }
+
+        try {
+            return InMemoryEmbeddingStore.fromFile(FALLBACK_EMB_FILE);
+        } catch (RuntimeException ignore) {
+            // Rebuild the embeddings
+        }
+
+        InMemoryEmbeddingStore<TextSegment> store = new InMemoryEmbeddingStore<>();
+        List<Document> documents = loader.load();
+        if (documents.isEmpty())
+            return store;
+        DocumentSplitter splitter = MarkdownCodeAwareSplitter.defaultSplitter();
+        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
+                .documentSplitter(splitter)
+                .embeddingModel(embeddingModel)
+                .embeddingStore(store)
+                .build();
+        documents.set(0, Document.from(
+                PerUserChatMaker.wrapUserInput("regenerateEmbeddings",
+                        documents.get(0).text())));
+        ingestor.ingest(documents);
+        store.serializeToFile(FALLBACK_EMB_FILE);
+
+        return store;
+    }
+
+    @Bean
+    ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore,
+                                      EmbeddingModel embeddingModel) {
+        return EmbeddingStoreContentRetriever.builder()
+                .embeddingStore(embeddingStore)
+                .embeddingModel(embeddingModel)
+                .maxResults(25)
+                .minScore(0.55)
+                .build();
+    }
+}

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/config/WebSocketConfig.java

@@ -74,6 +74,8 @@ public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
     public static final String IMPORT_QSMSG_TOPIC = IMPORT_TOPIC + "/qsmsg";
     public static final String INIT_IMPORT_QSMSG_TOPIC = INIT_IMPORT_TOPIC + "/qsmsg";
 
+    public static final String GENAI_QQL_TOPIC = TOPIC + "/genai-qql";
+
     public static final String SUBSCRIPTIONS_METRIC = "websocket.subscriptions";
 
     public static final String SEND_MESSAGES_METRIC = "websocket.messages";

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/controllers/GenAiController.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2025 EPAM Systems, Inc
+ *
+ * See the NOTICE file distributed with this work for additional information
+ * regarding copyright ownership. 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 com.epam.deltix.tbwg.webapp.controllers;
+
+import com.epam.deltix.tbwg.webapp.config.WebSocketConfig;
+import com.epam.deltix.tbwg.webapp.services.genai.GenAiService;
+import com.epam.deltix.tbwg.webapp.websockets.subscription.Subscription;
+import com.epam.deltix.tbwg.webapp.websockets.subscription.SubscriptionChannel;
+import com.epam.deltix.tbwg.webapp.websockets.subscription.SubscriptionController;
+import com.epam.deltix.tbwg.webapp.websockets.subscription.SubscriptionControllerRegistry;
+import org.jetbrains.annotations.Nullable;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.stereotype.Controller;
+import org.springframework.web.bind.annotation.CrossOrigin;
+
+import java.security.Principal;
+
+@Controller
+@CrossOrigin
+public class GenAiController implements SubscriptionController {
+
+    private final @Nullable GenAiService genAiService;
+
+    public GenAiController(SubscriptionControllerRegistry registry,
+                           @Autowired(required = false) @Nullable GenAiService genAiService) {
+        registry.register(WebSocketConfig.GENAI_QQL_TOPIC, this);
+        this.genAiService = genAiService;
+    }
+
+    @Override
+    public Subscription onSubscribe(SimpMessageHeaderAccessor headerAccessor, SubscriptionChannel channel) {
+        if (genAiService == null) {
+            channel.sendError("Gen AI service is disabled");
+            return () -> {};
+        }
+        String userInput = headerAccessor.getFirstNativeHeader("userInput");
+        String rawStreamKeys = headerAccessor.getFirstNativeHeader("streamKeys");
+        if (userInput == null || userInput.isEmpty()) {
+            channel.sendError(new IllegalArgumentException("userInput header is required"));
+            return () -> {};
+        }
+        Principal user = headerAccessor.getUser();
+        if (user == null) {
+            channel.sendError(new IllegalArgumentException("User header is required"));
+            return () -> {};
+        }
+        String username = user.getName();
+        if (username == null || username.isEmpty()) {
+            channel.sendError(new IllegalArgumentException("Username is required"));
+            return () -> {};
+        }
+
+        genAiService.subscribe(username, userInput, rawStreamKeys, channel);
+        return () -> genAiService.unsubscribe(channel);
+    }
+}

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/model/authorization/UserDto.java

@@ -31,6 +31,8 @@ public class UserDto {
     @NotEmpty
     private List<String> authorities;
 
+    private String aiApiKey;
+
     public String getUsername() {
         return username;
     }
@@ -54,4 +56,12 @@ public List<String> getAuthorities() {
     public void setAuthorities(List<String> authorities) {
         this.authorities = authorities;
     }
-}
+
+    public String getAiApiKey() {
+        return aiApiKey;
+    }
+
+    public void setAiApiKey(String aiApiKey) {
+        this.aiApiKey = aiApiKey;
+    }
+}

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/model/genai/QqlGenMessage.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2025 EPAM Systems, Inc
+ *
+ * See the NOTICE file distributed with this work for additional information
+ * regarding copyright ownership. 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 com.epam.deltix.tbwg.webapp.model.genai;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class QqlGenMessage {
+
+    @JsonProperty
+    private final String stage;
+    @JsonProperty
+    private final Integer attempt;
+    @JsonProperty
+    private final Integer maxAttempts;
+    @JsonProperty
+    private final String data;
+    @JsonProperty
+    private final String error;
+    @JsonProperty
+    private final Boolean finalEvent;
+
+    private QqlGenMessage(Builder b) {
+        this.stage = b.stage;
+        this.attempt = b.attempt;
+        this.maxAttempts = b.maxAttempts;
+        this.data = b.data;
+        this.error = b.error;
+        this.finalEvent = b.finalEvent;
+    }
+
+    public String getStage() { return stage; }
+    public Integer getAttempt() { return attempt; }
+    public Integer getMaxAttempts() { return maxAttempts; }
+    public String getData() { return data; }
+    public String getError() { return error; }
+    public Boolean getFinalEvent() { return finalEvent; }
+
+    public static Builder builder(String stage) { return new Builder(stage); }
+
+    public static class Builder {
+        private final String stage;
+        private Integer attempt;
+        private Integer maxAttempts;
+        private String data;
+        private String error;
+        private Boolean finalEvent;
+
+        public Builder(String stage) { this.stage = stage; }
+        public Builder attempt(Integer v) { this.attempt = v; return this; }
+        public Builder maxAttempts(Integer v) { this.maxAttempts = v; return this; }
+        public Builder data(String v) { this.data = v; return this; }
+        public Builder error(String v) { this.error = v; return this; }
+        public Builder finalEvent(Boolean v) { this.finalEvent = v; return this; }
+        public QqlGenMessage build() { return new QqlGenMessage(this); }
+    }
+}

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/services/authorization/FileAuthorizationProvider.java

@@ -97,7 +97,8 @@ private synchronized void reloadFile() {
                     users.put(
                         user.getUsername(),
                         new TbwgUser(
-                            user.getUsername(), user.getPassword(), buildAuthorities(user.getAuthorities())
+                            user.getUsername(), user.getPassword(), buildAuthorities(user.getAuthorities()),
+                                user.getAiApiKey()
                         )
                     );
                 });

java/ws-server/src/main/java/com/epam/deltix/tbwg/webapp/services/authorization/SettingsAuthorizationProvider.java

@@ -59,7 +59,8 @@ public SettingsAuthorizationProvider(SecurityOauth2ProviderSettings providerSett
                     new TbwgUser(
                         user.getUsername(),
                         providerType == ProviderType.BUILT_IN_OAUTH ? user.getPassword() : "",
-                        buildAuthorities(user.getAuthorities())
+                        buildAuthorities(user.getAuthorities()),
+                        user.getAiApiKey()
                     )
                 );
             });