nicolo-rinaldi commented on code in PR #4259: URL: https://github.com/apache/solr/pull/4259#discussion_r3138166558
########## solr/modules/language-models/src/java/org/apache/solr/languagemodels/documentenrichment/model/SolrFieldGenerationModel.java: ########## @@ -0,0 +1,256 @@ +/* + * 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.Capability; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.request.ResponseFormat; +import dev.langchain4j.model.chat.request.json.JsonArraySchema; +import dev.langchain4j.model.chat.request.json.JsonBooleanSchema; +import dev.langchain4j.model.chat.request.json.JsonIntegerSchema; +import dev.langchain4j.model.chat.request.json.JsonNumberSchema; +import dev.langchain4j.model.chat.request.json.JsonObjectSchema; +import dev.langchain4j.model.chat.request.json.JsonSchemaElement; +import dev.langchain4j.model.chat.request.json.JsonStringSchema; +import java.lang.invoke.MethodHandles; +import java.lang.reflect.Method; +import java.time.Duration; +import java.util.ArrayList; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +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.update.processor.DocumentEnrichmentUpdateProcessorFactory; +import org.apache.solr.languagemodels.documentenrichment.store.FieldGenerationModelException; +import org.apache.solr.languagemodels.documentenrichment.store.rest.ManagedFieldGenerationModelStore; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * This object wraps a {@link dev.langchain4j.model.chat.ChatModel} to produce the content of a + * field based on the content of other fields specified as input. It's meant to be used as a managed + * resource with the {@link ManagedFieldGenerationModelStore} + */ +public class SolrFieldGenerationModel implements Accountable { + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + private static final long BASE_RAM_BYTES = + RamUsageEstimator.shallowSizeOfInstance(SolrFieldGenerationModel.class); + private static final String TIMEOUT_PARAM = "timeout"; + 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 SolrFieldGenerationModel getInstance( + SolrResourceLoader solrResourceLoader, + String className, + String name, + Map<String, Object> params) + throws FieldGenerationModelException { + 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 chatModel; + 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}, which 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); + } + } + } + } + } + } + + // Always enforce strict schema adherence where supported. For Anthropic and Google it's enabled by default + if (!"dev.langchain4j.model.anthropic.AnthropicChatModel".equals(className) && + !"dev.langchain4j.model.googleai.GoogleAiGeminiChatModel".equals(className)) { + try { + builder + .getClass() + .getMethod("strictJsonSchema", Boolean.class) + .invoke(builder, true); + } catch (NoSuchMethodException ignored) { + log.debug("Model {} does not have strictJsonSchema param, structured output is not enforced", className); + } + } + chatModel = (ChatModel) builder.getClass().getMethod("build").invoke(builder); + return new SolrFieldGenerationModel(name, chatModel, params); + } catch (final Exception e) { + throw new FieldGenerationModelException("Model loading failed for " + className, e); + } + } + + public SolrFieldGenerationModel(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 Integer}, {@link + * Boolean}, or {@link java.util.List} depending on the Solr output field type + */ + public Object generateFieldValue(String prompt, ResponseFormat responseFormat) { + ChatRequest chatRequest = + ChatRequest.builder() + .responseFormat(responseFormat) + .messages(UserMessage.from(prompt)) + .build(); + String rawJson = chatModel.chat(chatRequest).aiMessage().text(); + Object parsed = Utils.fromJSONString(rawJson); + // It makes sense to keep this due to Ollama support Review Comment: Changed ########## solr/modules/language-models/src/test/org/apache/solr/languagemodels/documentenrichment/model/DummyChatModel.java: ########## @@ -0,0 +1,85 @@ +/* Review Comment: Yes, but this is a mock for the actual langchain4j ChatModel class, indeed it implements the ChatModel interface, so I thought the name must contain some reference to the langchain4j library. ########## solr/modules/language-models/src/test/org/apache/solr/languagemodels/documentenrichment/model/DummyChatModelTest.java: ########## @@ -0,0 +1,48 @@ +/* Review Comment: same as DummyChatModel ########## solr/modules/language-models/src/test-files/solr/collection1/conf/schema-language-models.xml: ########## @@ -36,11 +40,50 @@ <field name="2048_byte_vector" type="high_dimensional_byte_knn_vector" indexed="true" stored="true" /> <field name="2048_float_vector" type="high_dimensional_float_knn_vector" indexed="true" stored="true" /> <field name="string_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="body_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="tags_field" type="string" indexed="true" stored="true" multiValued="true" required="false"/> + <field name="enriched_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="enriched_field_multi" type="string" indexed="true" stored="true" multiValued="true" required="false"/> + + <!-- single-valued output fields for each supported type --> + <field name="output_long" type="plong" indexed="true" stored="true" multiValued="false"/> + <field name="output_int" type="pint" indexed="true" stored="true" multiValued="false"/> + <field name="output_float" type="pfloat" indexed="true" stored="true" multiValued="false"/> + <field name="output_double" type="pdouble" indexed="true" stored="true" multiValued="false"/> + <field name="output_boolean" type="boolean" indexed="true" stored="true" multiValued="false"/> + <field name="output_date" type="pdate" indexed="true" stored="true" multiValued="false"/> + + <!-- multivalued output fields for each supported type --> + <field name="output_long_multi" type="plong" indexed="true" stored="true" multiValued="true"/> + <field name="output_int_multi" type="pint" indexed="true" stored="true" multiValued="true"/> + <field name="output_float_multi" type="pfloat" indexed="true" stored="true" multiValued="true"/> + <field name="output_double_multi" type="pdouble" indexed="true" stored="true" multiValued="true"/> + <field name="output_boolean_multi" type="boolean" indexed="true" stored="true" multiValued="true"/> + <field name="output_date_multi" type="pdate" indexed="true" stored="true" multiValued="true"/> + <field name="output_binary" type="binary" indexed="false" stored="true" multiValued="false"/> + + <!-- output fields for unsupported types --> + <field name="output_bynary" type="binary" indexed="true" stored="true" multiValued="false"/> + + <!-- output fields for types explicitly unsupported (without they are supported via inheritance) --> Review Comment: Changed, see if now is clear ########## solr/modules/language-models/src/test/org/apache/solr/languagemodels/documentenrichment/model/ExceptionThrowingChatModel.java: ########## @@ -0,0 +1,48 @@ +/* Review Comment: same as DummyChatModel ########## solr/modules/language-models/src/test-files/solr/collection1/conf/schema-language-models.xml: ########## @@ -36,11 +40,50 @@ <field name="2048_byte_vector" type="high_dimensional_byte_knn_vector" indexed="true" stored="true" /> <field name="2048_float_vector" type="high_dimensional_float_knn_vector" indexed="true" stored="true" /> <field name="string_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="body_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="tags_field" type="string" indexed="true" stored="true" multiValued="true" required="false"/> + <field name="enriched_field" type="string" indexed="true" stored="true" multiValued="false" required="false"/> + <field name="enriched_field_multi" type="string" indexed="true" stored="true" multiValued="true" required="false"/> + + <!-- single-valued output fields for each supported type --> + <field name="output_long" type="plong" indexed="true" stored="true" multiValued="false"/> + <field name="output_int" type="pint" indexed="true" stored="true" multiValued="false"/> + <field name="output_float" type="pfloat" indexed="true" stored="true" multiValued="false"/> + <field name="output_double" type="pdouble" indexed="true" stored="true" multiValued="false"/> + <field name="output_boolean" type="boolean" indexed="true" stored="true" multiValued="false"/> + <field name="output_date" type="pdate" indexed="true" stored="true" multiValued="false"/> + + <!-- multivalued output fields for each supported type --> + <field name="output_long_multi" type="plong" indexed="true" stored="true" multiValued="true"/> + <field name="output_int_multi" type="pint" indexed="true" stored="true" multiValued="true"/> + <field name="output_float_multi" type="pfloat" indexed="true" stored="true" multiValued="true"/> + <field name="output_double_multi" type="pdouble" indexed="true" stored="true" multiValued="true"/> + <field name="output_boolean_multi" type="boolean" indexed="true" stored="true" multiValued="true"/> + <field name="output_date_multi" type="pdate" indexed="true" stored="true" multiValued="true"/> + <field name="output_binary" type="binary" indexed="false" stored="true" multiValued="false"/> + + <!-- output fields for unsupported types --> + <field name="output_bynary" type="binary" indexed="true" stored="true" multiValued="false"/> Review Comment: I didn't even noticed this. "output_bynary" is never used. Removed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
