Move inference result type check to the InferenceOperator

Author: afoucret

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

@@ -23,7 +23,7 @@
 
 import java.util.List;
 
-public class CompletionOperator extends InferenceOperator<Page> {
+public class CompletionOperator extends InferenceOperator<Page, ChatCompletionResults> {
 
     public record Factory(InferenceRunner inferenceRunner, String inferenceId, ExpressionEvaluator.Factory promptEvaluatorFactory)
         implements
@@ -48,7 +48,7 @@ public CompletionOperator(
         String inferenceId,
         ExpressionEvaluator promptEvaluator
     ) {
-        super(driverContext, inferenceRunner.getThreadContext(), inferenceRunner, inferenceId);
+        super(driverContext, inferenceRunner.getThreadContext(), inferenceRunner, inferenceId, ChatCompletionResults.class);
         this.promptEvaluator = promptEvaluator;
         this.blockFactory = driverContext.blockFactory();
     }
@@ -61,7 +61,7 @@ protected void performAsync(Page inputPage, ActionListener<Page> listener) {
         CountDownActionListener countDownListener = new CountDownActionListener(
             inputPage.getPositionCount(),
             listener.delegateFailureIgnoreResponseAndWrap(l -> {
-                try(BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(pageSize)) {
+                try (BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(pageSize)) {
                     BytesRefBuilder bytesRefBuilder = new BytesRefBuilder();
                     for (int pos = 0; pos < pageSize; pos++) {
                         if (responses[pos] == null) {
@@ -91,21 +91,12 @@ protected void performAsync(Page inputPage, ActionListener<Page> listener) {
                     }
 
                     InferenceAction.Request request = InferenceAction.Request.builder(inferenceId(), TaskType.COMPLETION)
-                        .setInput(List.of(promptBuilder.toString())).build();
+                        .setInput(List.of(promptBuilder.toString()))
+                        .build();
 
                     doInference(request, countDownListener.delegateFailureAndWrap((l, r) -> {
-                        if (r.getResults() instanceof ChatCompletionResults completionResults) {
-                            responses[currentPos] = completionResults.results().getFirst().content();
-                            l.onResponse(null);
-                        } else {
-                            l.onFailure(new IllegalStateException(
-                                "Inference result has wrong type. Got ["
-                                    + r.getResults().getClass()
-                                    + "] while expecting ["
-                                    + ChatCompletionResults.class
-                                    + "]"
-                            ));
-                        }
+                        responses[currentPos] = r.results().getFirst().content();
+                        l.onResponse(null);
                     }));
                 }
             }
@@ -131,5 +122,4 @@ public Page getOutput() {
     public String toString() {
         return "CompletionOperator[inference_id=[" + inferenceId() + "]]";
     }
-
 }

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

@@ -11,34 +11,52 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.compute.operator.AsyncOperator;
 import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 
-abstract public class InferenceOperator<Fetched> extends AsyncOperator<Fetched> {
+import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
+
+abstract public class InferenceOperator<Fetched, InferenceResult extends InferenceServiceResults> extends AsyncOperator<Fetched> {
 
     // Move to a setting.
     private static final int MAX_INFERENCE_WORKER = 10;
-
     private final InferenceRunner inferenceRunner;
     private final String inferenceId;
+    private final Class<InferenceResult> inferenceResultClass;
 
     public InferenceOperator(
         DriverContext driverContext,
         ThreadContext threadContext,
         InferenceRunner inferenceRunner,
-        String inferenceId
+        String inferenceId,
+        Class<InferenceResult> inferenceResultClass
     ) {
         super(driverContext, threadContext, MAX_INFERENCE_WORKER);
         this.inferenceRunner = inferenceRunner;
         this.inferenceId = inferenceId;
+        this.inferenceResultClass = inferenceResultClass;
 
         assert inferenceRunner.getThreadContext() != null;
     }
 
-    protected void doInference(InferenceAction.Request inferenceRequest, ActionListener<InferenceAction.Response> listener) {
-        inferenceRunner.doInference(inferenceRequest, listener);
+    protected final void doInference(InferenceAction.Request inferenceRequest, ActionListener<InferenceResult> listener) {
+        inferenceRunner.doInference(inferenceRequest, listener.map(this::checkedInferenceResults));
     }
 
     protected String inferenceId() {
         return inferenceId;
     }
+
+    private InferenceResult checkedInferenceResults(InferenceAction.Response inferenceResponse) {
+        if (inferenceResultClass.isInstance(inferenceResponse.getResults())) {
+            return inferenceResultClass.cast(inferenceResponse.getResults());
+        }
+        throw new IllegalStateException(
+            format(
+                "Inference result has wrong type. Got [{}] while expecting [{}]",
+                inferenceResponse.getResults().getClass().getName(),
+                inferenceResultClass.getName()
+            )
+        );
+    }
 }

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

@@ -25,7 +25,7 @@
 
 import java.util.List;
 
-public class RerankOperator extends InferenceOperator<Page> {
+public class RerankOperator extends InferenceOperator<Page, RankedDocsResults> {
     public record Factory(
         InferenceRunner inferenceRunner,
         String inferenceId,
@@ -65,7 +65,7 @@ public RerankOperator(
         ExpressionEvaluator rowEncoder,
         int scoreChannel
     ) {
-        super(driverContext, inferenceRunner.getThreadContext(), inferenceRunner, inferenceId);
+        super(driverContext, inferenceRunner.getThreadContext(), inferenceRunner, inferenceId, RankedDocsResults.class);
 
         this.blockFactory = driverContext.blockFactory();
         this.queryText = queryText;
@@ -81,10 +81,7 @@ protected void performAsync(Page inputPage, ActionListener<Page> listener) {
         try {
             doInference(
                 buildInferenceRequest(inputPage),
-                ActionListener.wrap(
-                    inferenceResponse -> outputListener.onResponse(buildOutput(inputPage, inferenceResponse)),
-                    outputListener::onFailure
-                )
+                outputListener.delegateFailureAndWrap((l, r) -> l.onResponse(buildOutput(inputPage, r)))
             );
         } catch (Exception e) {
             outputListener.onFailure(e);
@@ -111,21 +108,6 @@ public String toString() {
         return "RerankOperator[inference_id=[" + inferenceId() + "], query=[" + queryText + "], score_channel=[" + scoreChannel + "]]";
     }
 
-    private Page buildOutput(Page inputPage, InferenceAction.Response inferenceResponse) {
-        if (inferenceResponse.getResults() instanceof RankedDocsResults rankedDocsResults) {
-            return buildOutput(inputPage, rankedDocsResults);
-
-        }
-
-        throw new IllegalStateException(
-            "Inference result has wrong type. Got ["
-                + inferenceResponse.getResults().getClass()
-                + "] while expecting ["
-                + RankedDocsResults.class
-                + "]"
-        );
-    }
-
     private Page buildOutput(Page inputPage, RankedDocsResults rankedDocsResults) {
         int blockCount = Integer.max(inputPage.getBlockCount(), scoreChannel + 1);
         Block[] blocks = new Block[blockCount];