SOLR-18187: Document enrichment with LLMs

changelog/unreleased/llm-document-enrichment.yml

@@ -0,0 +1,4 @@
+title: Add DocumentEnrichmentUpdateProcessorFactory for LLM-based document enrichment at index time
+type: added # added, changed, fixed, deprecated, removed, dependency_update, security, other
+authors:
+- name: Nicolò Rinaldi

gradle/libs.versions.toml

@@ -449,11 +449,14 @@ ktor-client-serialization-json = { module = "io.ktor:ktor-serialization-kotlinx-
 ktor-server-cio = { module = "io.ktor:ktor-server-cio", version.ref = "ktor" }
 ktor-server-core = { module = "io.ktor:ktor-server-core", version.ref = "ktor" }
 ktor-server-htmlBuilder = { module = "io.ktor:ktor-server-html-builder", version.ref = "ktor" }
+langchain4j-anthropic = { module = "dev.langchain4j:langchain4j-anthropic" }
 langchain4j-bom = { module = "dev.langchain4j:langchain4j-bom", version.ref = "langchain4j-bom" }
 langchain4j-cohere = { module = "dev.langchain4j:langchain4j-cohere" }
 langchain4j-core = { module = "dev.langchain4j:langchain4j-core" }
+langchain4j-google-ai-gemini = { module = "dev.langchain4j:langchain4j-google-ai-gemini" }
 langchain4j-hugging-face = { module = "dev.langchain4j:langchain4j-hugging-face" }
 langchain4j-mistral-ai = { module = "dev.langchain4j:langchain4j-mistral-ai" }
+langchain4j-ollama = { module = "dev.langchain4j:langchain4j-ollama" }
 langchain4j-open-ai = { module = "dev.langchain4j:langchain4j-open-ai" }
 lmax-disruptor = { module = "com.lmax:disruptor", version.ref = "lmax-disruptor" }
 locationtech-spatial4j = { module = "org.locationtech.spatial4j:spatial4j", version.ref = "spatial4j" }

solr/modules/language-models/build.gradle

@@ -29,9 +29,12 @@ dependencies {
   implementation libs.apache.lucene.core
 
   implementation libs.langchain4j.core
+  runtimeOnly libs.langchain4j.anthropic
   runtimeOnly libs.langchain4j.cohere
+  runtimeOnly libs.langchain4j.google.ai.gemini
   runtimeOnly libs.langchain4j.hugging.face
   runtimeOnly libs.langchain4j.mistral.ai
+  runtimeOnly libs.langchain4j.ollama
   runtimeOnly libs.langchain4j.open.ai
 
   implementation libs.slf4j.api

solr/modules/language-models/gradle.lockfile

@@ -40,13 +40,16 @@ com.tdunning:t-digest:3.3=jarValidation,runtimeClasspath,runtimeLibs,solrPlatfor
 commons-cli:commons-cli:1.10.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
 commons-codec:commons-codec:1.19.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
 commons-io:commons-io:2.20.0=compileClasspath,jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testCompileClasspath,testRuntimeClasspath
+dev.langchain4j:langchain4j-anthropic:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-bom:1.9.1=compileClasspath,jarValidation,runtimeClasspath,runtimeLibs,testCompileClasspath,testRuntimeClasspath
 dev.langchain4j:langchain4j-cohere:1.9.1-beta17=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-core:1.9.1=compileClasspath,jarValidation,runtimeClasspath,runtimeLibs,testCompileClasspath,testRuntimeClasspath
+dev.langchain4j:langchain4j-google-ai-gemini:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-http-client-jdk:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-http-client:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-hugging-face:1.9.1-beta17=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-mistral-ai:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
+dev.langchain4j:langchain4j-ollama:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 dev.langchain4j:langchain4j-open-ai:1.9.1=jarValidation,runtimeClasspath,runtimeLibs,testRuntimeClasspath
 io.dropwizard.metrics:metrics-annotation:4.2.26=jarValidation,testRuntimeClasspath
 io.dropwizard.metrics:metrics-core:4.2.26=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath

solr/modules/language-models/src/java/org/apache/solr/languagemodels/documentenrichment/model/SolrChatModel.java

@@ -0,0 +1,215 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.apache.solr.languagemodels.documentenrichment.model;
+
+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 java.lang.invoke.MethodHandles;
+import java.lang.reflect.Method;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Map;
+import java.util.Objects;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
+import org.apache.solr.common.SolrException;
+import org.apache.solr.common.util.Utils;
+import org.apache.solr.core.SolrResourceLoader;
+import org.apache.solr.languagemodels.documentenrichment.store.ChatModelException;
+import org.apache.solr.languagemodels.documentenrichment.store.rest.ManagedChatModelStore;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This object wraps a {@link ChatModel} to produce the content of new fields from another.
+ * It's meant to be used as a managed resource with the {@link
+ * ManagedChatModelStore}
+ */
+public class SolrChatModel implements Accountable {
+  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+  private static final long BASE_RAM_BYTES =
+      RamUsageEstimator.shallowSizeOfInstance(SolrChatModel.class);
+  // timeout is type Duration
+  private static final String TIMEOUT_PARAM = "timeout";
+
+  // the followings are Integer type
+  private static final String MAX_RETRIES_PARAM = "maxRetries";
+  private static final String THINKING_BUDGET_TOKENS = "thinkingBudgetTokens";
+  private static final String RANDOM_SEED = "randomSeed";
+
+  private final String name;
+  private final Map<String, Object> params;
+  private final ChatModel chatModel;
+  private final int hashCode;
+
+  public static SolrChatModel getInstance(
+      SolrResourceLoader solrResourceLoader,
+      String className,
+      String name,
+      Map<String, Object> params)
+      throws ChatModelException {
+    try {
+      /*
+       * The idea here is to build a {@link dev.langchain4j.model.chat.ChatModel} using inversion
+       * of control.
+       * Each model has its own list of parameters we don't know beforehand, but each {@link dev.langchain4j.model.chat.ChatModel} class
+       * has its own builder that uses setters with the same name of the parameter in input.
+       * */
+      ChatModel textToTextModel;
+      Class<?> modelClass = solrResourceLoader.findClass(className, ChatModel.class);
+      var builder = modelClass.getMethod("builder").invoke(null);
+      if (params != null) {
+        /*
+         * This block of code has the responsibility of instantiate a {@link
+         * dev.langchain4j.model.chat.ChatModel} using the params provided.classes have
+         * params of The specific implementation of {@link
+         * dev.langchain4j.model.chat.ChatModel} is not known beforehand. So we benefit of
+         * the design choice in langchain4j that each subclass implementing {@link
+         * dev.langchain4j.model.chat.ChatModel} uses setters with the same name of the
+         * param.
+         */
+        for (String paramName : params.keySet()) {
+          /*
+           * When a param is not primitive, we need to instantiate the object explicitly and then call the
+           * setter method.
+           * N.B. when adding support to new models, pay attention to all the parameters they
+           * support, some of them may require to be handled in here as separate switch cases
+           */
+          switch (paramName) {
+            case TIMEOUT_PARAM -> builder
+                  .getClass()
+                  .getMethod(paramName, Duration.class)
+                  .invoke(builder, Duration.ofSeconds((Long) params.get(paramName)));
+
+            case MAX_RETRIES_PARAM, THINKING_BUDGET_TOKENS, RANDOM_SEED -> builder
+                .getClass()
+                .getMethod(paramName, Integer.class)
+                .invoke(builder, ((Long) params.get(paramName)).intValue());
+
+              /*
+               * For primitive params if there's only one setter available, we call it.
+               * If there's choice we default to the string one
+               */
+            default -> {
+              ArrayList<Method> paramNameMatches = new ArrayList<>();
+              for (var method : builder.getClass().getMethods()) {
+                if (paramName.equals(method.getName()) && method.getParameterCount() == 1) {
+                  paramNameMatches.add(method);
+                }
+              }
+              if (paramNameMatches.size() == 1) {
+                paramNameMatches.getFirst().invoke(builder, params.get(paramName));
+              } else {
+                try {
+                  builder
+                      .getClass()
+                      .getMethod(paramName, String.class)
+                      .invoke(builder, params.get(paramName).toString());
+                } catch (NoSuchMethodException e) {
+                  log.error("Parameter {} not supported by model {}", paramName, className);
+                  throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, e.getMessage(), e);
+                }
+              }
+            }
+          }
+        }
+      }
+      textToTextModel = (ChatModel) builder.getClass().getMethod("build").invoke(builder);
+      return new SolrChatModel(name, textToTextModel, params);
+    } catch (final Exception e) {
+      throw new ChatModelException("Model loading failed for " + className, e);
+    }
+  }
+
+  public SolrChatModel(
+      String name, ChatModel chatModel, Map<String, Object> params) {
+    this.name = name;
+    this.chatModel = chatModel;
+    this.params = params;
+    this.hashCode = calculateHashCode();
+  }
+
+  /**
+   * Sends a structured chat request and returns the parsed value from the {@code {"value": ...}}
+   * JSON object that the model is instructed to produce via {@code responseFormat}.
+   *
+   * @return the extracted value: a {@link String}, {@link Number}, {@link Boolean}, or {@link
+   *     java.util.List} depending on the Solr output field type
+   */
+  public Object chat(String text, ResponseFormat responseFormat) {
+    ChatRequest chatRequest =
+        ChatRequest.builder()
+            .responseFormat(responseFormat)
+            .messages(UserMessage.from(text))
+            .build();
+    String rawJson = chatModel.chat(chatRequest).aiMessage().text();
+    Object parsed = Utils.fromJSONString(rawJson);
+    if (!(parsed instanceof Map<?, ?> map) || !map.containsKey("value")) {
+      throw new SolrException(
+          SolrException.ErrorCode.SERVER_ERROR,
+          "LLM response is missing the 'value' key: " + rawJson);
+    }
+    return map.get("value");
+  }
+
+  @Override
+  public String toString() {
+    return getClass().getSimpleName() + "(name=" + getName() + ")";
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return BASE_RAM_BYTES
+        + RamUsageEstimator.sizeOfObject(name)
+        + RamUsageEstimator.sizeOfObject(chatModel);
+  }
+
+  @Override
+  public int hashCode() {
+    return hashCode;
+  }
+
+  private int calculateHashCode() {
+    final int prime = 31;
+    int result = 1;
+    result = (prime * result) + Objects.hashCode(name);
+    result = (prime * result) + Objects.hashCode(chatModel);
+    return result;
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) return true;
+    if (!(obj instanceof SolrChatModel)) return false;
+    final SolrChatModel other = (SolrChatModel) obj;
+    return Objects.equals(chatModel, other.chatModel) && Objects.equals(name, other.name);
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public String getChatModelClassName() {
+    return chatModel.getClass().getName();
+  }
+
+  public Map<String, Object> getParams() {
+    return params;
+  }
+}

solr/modules/language-models/src/java/org/apache/solr/languagemodels/documentenrichment/model/package-info.java

@@ -0,0 +1,19 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+
+/** APIs and classes for implementing text to vector logic. */
+package org.apache.solr.languagemodels.documentenrichment.model;

solr/modules/language-models/src/java/org/apache/solr/languagemodels/documentenrichment/store/ChatModelException.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.apache.solr.languagemodels.documentenrichment.store;
+
+public class ChatModelException extends RuntimeException {
+
+  private static final long serialVersionUID = 1L;
+
+  public ChatModelException(String message) {
+    super(message);
+  }
+
+  public ChatModelException(String message, Exception cause) {
+    super(message, cause);
+  }
+}

solr/modules/language-models/src/java/org/apache/solr/languagemodels/documentenrichment/store/ChatModelStore.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.apache.solr.languagemodels.documentenrichment.store;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.solr.languagemodels.documentenrichment.model.SolrChatModel;
+
+/** Simple store to manage CRUD operations on the {@link SolrChatModel} */
+public class ChatModelStore {
+
+  private final Map<String, SolrChatModel> availableModels;
+
+  public ChatModelStore() {
+    availableModels = Collections.synchronizedMap(new LinkedHashMap<>());
+  }
+
+  public SolrChatModel getModel(String name) {
+    return availableModels.get(name);
+  }
+
+  public void clear() {
+    availableModels.clear();
+  }
+
+  public List<SolrChatModel> getModels() {
+    synchronized (availableModels) {
+      final List<SolrChatModel> availableModelsValues =
+          new ArrayList<>(availableModels.values());
+      return Collections.unmodifiableList(availableModelsValues);
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "ChatModelStore [availableModels=" + availableModels.keySet() + "]";
+  }
+
+  public SolrChatModel delete(String modelName) {
+    return availableModels.remove(modelName);
+  }
+
+  public void addModel(SolrChatModel modeldata) throws ChatModelException {
+    final String name = modeldata.getName();
+    if (availableModels.putIfAbsent(modeldata.getName(), modeldata) != null) {
+      throw new ChatModelException(
+          "model '" + name + "' already exists. Please use a different name");
+    }
+  }
+}