Set model registry for each semantic query builder instance

Author: Mikep86

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/search/ccs/SemanticCrossClusterSearchIT.java

@@ -37,6 +37,7 @@
 import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
 import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
 import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
+import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.io.IOException;
 import java.util.Collection;
@@ -104,8 +105,12 @@ public void testSemanticCrossClusterSearch() throws Exception {
         String localIndex = (String) testClusterInfo.get("local.index");
         String remoteIndex = (String) testClusterInfo.get("remote.index");
 
+        ModelRegistry modelRegistry = cluster(LOCAL_CLUSTER).getCurrentMasterNodeInstance(ModelRegistry.class);
+        SemanticQueryBuilder queryBuilder = new SemanticQueryBuilder(INFERENCE_FIELD, "foo");
+        queryBuilder.setModelRegistrySupplier(() -> modelRegistry);
+
         SearchRequest searchRequest = new SearchRequest(localIndex, REMOTE_CLUSTER + ":" + remoteIndex);
-        searchRequest.source(new SearchSourceBuilder().query(new SemanticQueryBuilder(INFERENCE_FIELD, "foo")).size(10));
+        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).size(10));
         searchRequest.setCcsMinimizeRoundtrips(true);
 
         assertResponse(client(LOCAL_CLUSTER).search(searchRequest), response -> {
@@ -155,12 +160,7 @@ public void testMatchCrossClusterSearch() throws Exception {
     private Map<String, Object> setupTwoClusters(String[] localIndices, String[] remoteIndices) throws IOException {
         final String localInferenceId = "local_inference_id";
         final String remoteInferenceId = "remote_inference_id";
-
-        // TODO: Resolve bug where remote model registry overwrites local model registry in SemanticQueryBuilder
         createInferenceEndpoint(client(LOCAL_CLUSTER), TaskType.TEXT_EMBEDDING, localInferenceId, TEXT_EMBEDDING_SERVICE_SETTINGS_1);
-        createInferenceEndpoint(client(LOCAL_CLUSTER), TaskType.TEXT_EMBEDDING, remoteInferenceId, TEXT_EMBEDDING_SERVICE_SETTINGS_2);
-
-        createInferenceEndpoint(client(REMOTE_CLUSTER), TaskType.TEXT_EMBEDDING, localInferenceId, TEXT_EMBEDDING_SERVICE_SETTINGS_1);
         createInferenceEndpoint(client(REMOTE_CLUSTER), TaskType.TEXT_EMBEDDING, remoteInferenceId, TEXT_EMBEDDING_SERVICE_SETTINGS_2);
 
         int numShardsLocal = randomIntBetween(2, 10);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -568,7 +568,15 @@ public Collection<MappedActionFilter> getMappedActionFilters() {
     }
 
     public List<QuerySpec<?>> getQueries() {
-        return List.of(new QuerySpec<>(SemanticQueryBuilder.NAME, SemanticQueryBuilder::new, SemanticQueryBuilder::fromXContent));
+        return List.of(new QuerySpec<>(SemanticQueryBuilder.NAME, i -> {
+            SemanticQueryBuilder queryBuilder = new SemanticQueryBuilder(i);
+            queryBuilder.setModelRegistrySupplier(getModelRegistry());
+            return queryBuilder;
+        }, p -> {
+            SemanticQueryBuilder queryBuilder = SemanticQueryBuilder.fromXContent(p);
+            queryBuilder.setModelRegistrySupplier(getModelRegistry());
+            return queryBuilder;
+        }));
     }
 
     @Override
@@ -602,8 +610,6 @@ public void onNodeStarted() {
         if (registry != null) {
             registry.onNodeStarted();
         }
-
-        SemanticQueryBuilder.setModelRegistrySupplier(getModelRegistry());
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java

@@ -74,18 +74,14 @@ public class SemanticQueryBuilder extends AbstractQueryBuilder<SemanticQueryBuil
         declareStandardFields(PARSER);
     }
 
-    private static Supplier<ModelRegistry> MODEL_REGISTRY_SUPPLIER = () -> null;
-
-    public static void setModelRegistrySupplier(Supplier<ModelRegistry> supplier) {
-        MODEL_REGISTRY_SUPPLIER = supplier;
-    }
-
     private final String fieldName;
     private final String query;
     private final EmbeddingsProvider embeddingsProvider;
     private final boolean noInferenceResults;
     private final Boolean lenient;
 
+    private Supplier<ModelRegistry> modelRegistrySupplier = () -> null;
+
     public SemanticQueryBuilder(String fieldName, String query) {
         this(fieldName, query, null);
     }
@@ -126,6 +122,10 @@ public SemanticQueryBuilder(StreamInput in) throws IOException {
         }
     }
 
+    public void setModelRegistrySupplier(Supplier<ModelRegistry> supplier) {
+        modelRegistrySupplier = supplier;
+    }
+
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         out.writeString(fieldName);
@@ -150,6 +150,7 @@ private SemanticQueryBuilder(SemanticQueryBuilder other, EmbeddingsProvider embe
         this.embeddingsProvider = embeddingsProvider;
         this.noInferenceResults = noInferenceResults;
         this.lenient = other.lenient;
+        this.modelRegistrySupplier = other.modelRegistrySupplier;
     }
 
     @Override
@@ -208,7 +209,7 @@ private QueryBuilder doRewriteBuildSemanticQuery(SearchExecutionContext searchEx
                 );
             }
 
-            ModelRegistry modelRegistry = MODEL_REGISTRY_SUPPLIER.get();
+            ModelRegistry modelRegistry = modelRegistrySupplier.get();
             if (modelRegistry == null) {
                 throw new IllegalStateException("Model registry has not been set");
             }
@@ -272,7 +273,7 @@ private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext qu
 
         boolean modified = false;
         if (queryRewriteContext.hasAsyncActions() == false) {
-            ModelRegistry modelRegistry = MODEL_REGISTRY_SUPPLIER.get();
+            ModelRegistry modelRegistry = modelRegistrySupplier.get();
             if (modelRegistry == null) {
                 throw new IllegalStateException("Model registry has not been set");
             }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/LocalStateInferencePlugin.java

@@ -73,9 +73,4 @@ public Map<String, Highlighter> getHighlighters() {
     public Collection<MappedActionFilter> getMappedActionFilters() {
         return inferencePlugin.getMappedActionFilters();
     }
-
-    @Override
-    public void onNodeStarted() {
-        inferencePlugin.onNodeStarted();
-    }
 }