Add embedding model

Author: andreadimaio

bam/deployment/src/main/java/io/quarkiverse/langchain4j/bam/deployment/BamProcessor.java

@@ -1,6 +1,7 @@
 package io.quarkiverse.langchain4j.bam.deployment;
 
 import static io.quarkiverse.langchain4j.deployment.Langchain4jDotNames.CHAT_MODEL;
+import static io.quarkiverse.langchain4j.deployment.Langchain4jDotNames.EMBEDDING_MODEL;
 
 import java.util.Optional;
 
@@ -9,7 +10,9 @@
 import io.quarkiverse.langchain4j.bam.runtime.BamRecorder;
 import io.quarkiverse.langchain4j.bam.runtime.config.Langchain4jBamConfig;
 import io.quarkiverse.langchain4j.deployment.items.ChatModelProviderCandidateBuildItem;
+import io.quarkiverse.langchain4j.deployment.items.EmbeddingModelProviderCandidateBuildItem;
 import io.quarkiverse.langchain4j.deployment.items.SelectedChatModelProviderBuildItem;
+import io.quarkiverse.langchain4j.deployment.items.SelectedEmbeddingModelCandidateBuildItem;
 import io.quarkus.arc.deployment.SyntheticBeanBuildItem;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
@@ -30,19 +33,27 @@ FeatureBuildItem feature() {
 
     @BuildStep
     public void providerCandidates(BuildProducer<ChatModelProviderCandidateBuildItem> chatProducer,
+            BuildProducer<EmbeddingModelProviderCandidateBuildItem> embeddingProducer,
             Langchain4jBamBuildConfig config) {
+
         if (config.chatModel().enabled().isEmpty() || config.chatModel().enabled().get()) {
             chatProducer.produce(new ChatModelProviderCandidateBuildItem(PROVIDER));
         }
+
+        if (config.embeddingModel().enabled().isEmpty() || config.embeddingModel().enabled().get()) {
+            embeddingProducer.produce(new EmbeddingModelProviderCandidateBuildItem(PROVIDER));
+        }
     }
 
     @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
     @BuildStep
     @Record(ExecutionTime.RUNTIME_INIT)
     void generateBeans(BamRecorder recorder,
             Optional<SelectedChatModelProviderBuildItem> selectedChatItem,
+            Optional<SelectedEmbeddingModelCandidateBuildItem> selectedEmbedding,
             Langchain4jBamConfig config,
             BuildProducer<SyntheticBeanBuildItem> beanProducer) {
+
         if (selectedChatItem.isPresent() && PROVIDER.equals(selectedChatItem.get().getProvider())) {
             beanProducer.produce(SyntheticBeanBuildItem
                     .configure(CHAT_MODEL)
@@ -52,5 +63,17 @@ void generateBeans(BamRecorder recorder,
                     .supplier(recorder.chatModel(config))
                     .done());
         }
+
+        if (selectedEmbedding.isPresent() && PROVIDER.equals(selectedEmbedding.get().getProvider())) {
+            beanProducer.produce(
+                    SyntheticBeanBuildItem
+                            .configure(EMBEDDING_MODEL)
+                            .setRuntimeInit()
+                            .defaultBean()
+                            .scope(ApplicationScoped.class)
+                            .supplier(recorder.embeddingModel(config))
+                            .unremovable()
+                            .done());
+        }
     }
 }

bam/deployment/src/main/java/io/quarkiverse/langchain4j/bam/deployment/ChatModelBuildConfig.java

@@ -13,4 +13,9 @@ public interface ChatModelBuildConfig {
      */
     @ConfigDocDefault("true")
     Optional<Boolean> enabled();
+
+    /**
+     * Embedding model related settings
+     */
+    EmbeddingModelBuildConfig embeddingModel();
 }

bam/deployment/src/main/java/io/quarkiverse/langchain4j/bam/deployment/EmbeddingModelBuildConfig.java

@@ -0,0 +1,16 @@
+package io.quarkiverse.langchain4j.bam.deployment;
+
+import java.util.Optional;
+
+import io.quarkus.runtime.annotations.ConfigDocDefault;
+import io.quarkus.runtime.annotations.ConfigGroup;
+
+@ConfigGroup
+public interface EmbeddingModelBuildConfig {
+
+    /**
+     * Whether the model should be enabled
+     */
+    @ConfigDocDefault("true")
+    Optional<Boolean> enabled();
+}

bam/deployment/src/main/java/io/quarkiverse/langchain4j/bam/deployment/Langchain4jBamBuildConfig.java

@@ -13,4 +13,9 @@ public interface Langchain4jBamBuildConfig {
      * Chat model related settings
      */
     ChatModelBuildConfig chatModel();
+
+    /**
+     * Embedding model related settings
+     */
+    EmbeddingModelBuildConfig embeddingModel();
 }

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/BamEmbeddingModel.java

@@ -0,0 +1,129 @@
+package io.quarkiverse.langchain4j.bam;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.net.URL;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+import org.jboss.resteasy.reactive.client.api.LoggingScope;
+
+import dev.langchain4j.data.embedding.Embedding;
+import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.model.embedding.EmbeddingModel;
+import dev.langchain4j.model.embedding.TokenCountEstimator;
+import dev.langchain4j.model.output.Response;
+import io.quarkus.rest.client.reactive.QuarkusRestClientBuilder;
+
+public class BamEmbeddingModel implements EmbeddingModel, TokenCountEstimator {
+
+    private final String token;
+    private final String modelId;
+    private final String version;
+    public boolean logResponses;
+    public boolean logRequests;
+    private final BamRestApi client;
+
+    public BamEmbeddingModel(Builder config) {
+
+        QuarkusRestClientBuilder builder = QuarkusRestClientBuilder.newBuilder()
+                .baseUri(config.url)
+                .connectTimeout(config.timeout.toSeconds(), TimeUnit.SECONDS)
+                .readTimeout(config.timeout.toSeconds(), TimeUnit.SECONDS);
+
+        if (config.logRequests || config.logResponses) {
+            builder.loggingScope(LoggingScope.REQUEST_RESPONSE);
+            builder.clientLogger(new BamRestApi.WatsonClientLogger(
+                    config.logRequests,
+                    config.logResponses));
+        }
+
+        this.client = builder.build(BamRestApi.class);
+        this.token = config.accessToken;
+        this.modelId = config.modelId;
+        this.version = config.version;
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    @Override
+    public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
+
+        List<Embedding> result = new ArrayList<>();
+        for (TextSegment textSegment : textSegments) {
+
+            var request = new EmbeddingRequest(modelId, textSegment.text());
+            var response = client.embeddings(request, token, version);
+
+            var vector = response.results().get(0);
+            result.add(Embedding.from(vector));
+        }
+
+        return Response.from(result);
+    }
+
+    @Override
+    public int estimateTokenCount(String text) {
+
+        var request = new TokenizationRequest(modelId, text);
+        return client.tokenization(request, token, version).tokenCount();
+    }
+
+    public static final class Builder {
+
+        private String accessToken;
+        private String version;
+        private URI url = URI.create("https://bam-api.res.ibm.com");
+        private Duration timeout = Duration.ofSeconds(15);
+        private String modelId;
+        public boolean logResponses;
+        public boolean logRequests;
+
+        public Builder accessToken(String accessToken) {
+            this.accessToken = accessToken;
+            return this;
+        }
+
+        public Builder version(String version) {
+            this.version = version;
+            return this;
+        }
+
+        public Builder url(URL url) {
+            try {
+                this.url = url.toURI();
+            } catch (URISyntaxException e) {
+                throw new RuntimeException(e);
+            }
+            return this;
+        }
+
+        public Builder timeout(Duration timeout) {
+            this.timeout = timeout;
+            return this;
+        }
+
+        public Builder modelId(String modelId) {
+            this.modelId = modelId;
+            return this;
+        }
+
+        public Builder logRequests(boolean logRequests) {
+            this.logRequests = logRequests;
+            return this;
+        }
+
+        public Builder logResponses(boolean logResponses) {
+            this.logResponses = logResponses;
+            return this;
+        }
+
+        public BamEmbeddingModel build() {
+            return new BamEmbeddingModel(this);
+        }
+    }
+}

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/BamRestApi.java

@@ -3,6 +3,7 @@
 import static java.util.stream.Collectors.joining;
 import static java.util.stream.StreamSupport.stream;
 
+import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -46,6 +47,10 @@ public interface BamRestApi {
     @Path("text/chat")
     TextGenerationResponse chat(TextGenerationRequest request, @NotBody String token, @QueryParam("version") String version);
 
+    @POST
+    @Path("/text/embeddings")
+    EmbeddingResponse embeddings(EmbeddingRequest request, @NotBody String token, @QueryParam("version") String version);
+
     @POST
     @Path("/text/tokenization")
     public TokenizationResponse tokenization(TokenizationRequest request, @NotBody String token,

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/EmbeddingRequest.java

@@ -0,0 +1,5 @@
+package io.quarkiverse.langchain4j.bam;
+
+public record EmbeddingRequest(String modelId, String input) {
+
+}

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/EmbeddingResponse.java

@@ -0,0 +1,7 @@
+package io.quarkiverse.langchain4j.bam;
+
+import java.util.List;
+
+public record EmbeddingResponse(List<List<Float>> results) {
+
+}

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/runtime/BamRecorder.java

@@ -3,6 +3,7 @@
 import static io.quarkiverse.langchain4j.runtime.OptionalUtil.firstOrDefault;
 import java.util.function.Supplier;
 import io.quarkiverse.langchain4j.bam.BamChatModel;
+import io.quarkiverse.langchain4j.bam.BamEmbeddingModel;
 import io.quarkiverse.langchain4j.bam.runtime.config.ChatModelConfig;
 import io.quarkiverse.langchain4j.bam.runtime.config.Langchain4jBamConfig;
 import io.quarkus.runtime.annotations.Recorder;
@@ -46,4 +47,26 @@ public Object get() {
             }
         };
     }
+
+    public Supplier<?> embeddingModel(Langchain4jBamConfig runtimeConfig) {
+
+        var embeddingModelConfig = runtimeConfig.embeddingModel();
+
+        var builder = BamEmbeddingModel.builder()
+                .accessToken(runtimeConfig.apiKey())
+                .timeout(runtimeConfig.timeout())
+                .version(runtimeConfig.version())
+                .modelId(embeddingModelConfig.modelId());
+
+        if (runtimeConfig.baseUrl().isPresent()) {
+            builder.url(runtimeConfig.baseUrl().get());
+        }
+
+        return new Supplier<>() {
+            @Override
+            public Object get() {
+                return builder.build();
+            }
+        };
+    }
 }

bam/runtime/src/main/java/io/quarkiverse/langchain4j/bam/runtime/config/EmbeddingModelConfig.java

@@ -0,0 +1,14 @@
+package io.quarkiverse.langchain4j.bam.runtime.config;
+
+import io.quarkus.runtime.annotations.ConfigGroup;
+import io.smallrye.config.WithDefault;
+
+@ConfigGroup
+public interface EmbeddingModelConfig {
+
+    /**
+     * Model to use
+     */
+    @WithDefault("ibm/slate.30m.english.rtrvr-26.10.2023")
+    String modelId();
+}