Add support for chat_completion to the Completion physical plan.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Completion.java

@@ -86,7 +86,6 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_CHAT_COMPLETION_SUPPORT)) {
             out.writeOptional((output, taskType) -> output.writeString(taskType.toString()), taskType());
         }
-
         out.writeNamedWriteable(prompt);
         out.writeNamedWriteable(targetField);
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/CompletionExec.java

@@ -7,9 +7,11 @@
 
 package org.elasticsearch.xpack.esql.plan.physical.inference;
 
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
@@ -37,8 +39,15 @@ public class CompletionExec extends InferenceExec {
     private final Attribute targetField;
     private List<Attribute> lazyOutput;
 
-    public CompletionExec(Source source, PhysicalPlan child, Expression inferenceId, Expression prompt, Attribute targetField) {
-        super(source, child, inferenceId);
+    public CompletionExec(
+        Source source,
+        PhysicalPlan child,
+        Expression inferenceId,
+        TaskType taskType,
+        Expression prompt,
+        Attribute targetField
+    ) {
+        super(source, child, inferenceId, taskType);
         this.prompt = prompt;
         this.targetField = targetField;
     }
@@ -48,6 +57,9 @@ public CompletionExec(StreamInput in) throws IOException {
             Source.readFrom((PlanStreamInput) in),
             in.readNamedWriteable(PhysicalPlan.class),
             in.readNamedWriteable(Expression.class),
+            in.getTransportVersion().onOrAfter(TransportVersions.ESQL_CHAT_COMPLETION_SUPPORT)
+                ? TaskType.fromString(in.readString())
+                : TaskType.COMPLETION,
             in.readNamedWriteable(Expression.class),
             in.readNamedWriteable(Attribute.class)
         );
@@ -61,6 +73,9 @@ public String getWriteableName() {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_CHAT_COMPLETION_SUPPORT)) {
+            out.writeString(taskType().toString());
+        }
         out.writeNamedWriteable(prompt);
         out.writeNamedWriteable(targetField);
     }
@@ -75,12 +90,12 @@ public Attribute targetField() {
 
     @Override
     protected NodeInfo<? extends PhysicalPlan> info() {
-        return NodeInfo.create(this, CompletionExec::new, child(), inferenceId(), prompt, targetField);
+        return NodeInfo.create(this, CompletionExec::new, child(), inferenceId(), taskType(), prompt, targetField);
     }
 
     @Override
     public UnaryExec replaceChild(PhysicalPlan newChild) {
-        return new CompletionExec(source(), newChild, inferenceId(), prompt, targetField);
+        return new CompletionExec(source(), newChild, inferenceId(), taskType(), prompt, targetField);
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/InferenceExec.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.plan.physical.inference;
 
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
@@ -18,16 +19,22 @@
 
 public abstract class InferenceExec extends UnaryExec {
     private final Expression inferenceId;
+    private final TaskType taskType;
 
-    protected InferenceExec(Source source, PhysicalPlan child, Expression inferenceId) {
+    protected InferenceExec(Source source, PhysicalPlan child, Expression inferenceId, TaskType taskType) {
         super(source, child);
         this.inferenceId = inferenceId;
+        this.taskType = taskType;
     }
 
     public Expression inferenceId() {
         return inferenceId;
     }
 
+    public TaskType taskType() {
+        return taskType;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         source().writeTo(out);
@@ -41,11 +48,11 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) return false;
         if (super.equals(o) == false) return false;
         InferenceExec that = (InferenceExec) o;
-        return inferenceId.equals(that.inferenceId);
+        return inferenceId.equals(that.inferenceId) && taskType == that.taskType;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), inferenceId());
+        return Objects.hash(super.hashCode(), inferenceId(), taskType);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/RerankExec.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
@@ -48,7 +49,7 @@ public RerankExec(
         List<Alias> rerankFields,
         Attribute scoreAttribute
     ) {
-        super(source, child, inferenceId);
+        super(source, child, inferenceId, TaskType.RERANK);
         this.queryText = queryText;
         this.rerankFields = rerankFields;
         this.scoreAttribute = scoreAttribute;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java

@@ -103,7 +103,14 @@ static PhysicalPlan mapUnary(UnaryPlan p, PhysicalPlan child) {
         }
 
         if (p instanceof Completion completion) {
-            return new CompletionExec(completion.source(), child, completion.inferenceId(), completion.prompt(), completion.targetField());
+            return new CompletionExec(
+                completion.source(),
+                child,
+                completion.inferenceId(),
+                completion.taskType(),
+                completion.prompt(),
+                completion.targetField()
+            );
         }
 
         if (p instanceof Enrich enrich) {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/inference/CompletionExecSerializationTests.java

@@ -7,11 +7,13 @@
 
 package org.elasticsearch.xpack.esql.plan.physical.inference;
 
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.expression.function.ReferenceAttributeTests;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.physical.AbstractPhysicalPlanSerializationTests;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
@@ -20,7 +22,14 @@
 public class CompletionExecSerializationTests extends AbstractPhysicalPlanSerializationTests<CompletionExec> {
     @Override
     protected CompletionExec createTestInstance() {
-        return new CompletionExec(randomSource(), randomChild(0), randomInferenceId(), randomPrompt(), randomAttribute());
+        return new CompletionExec(
+            randomSource(),
+            randomChild(0),
+            randomInferenceId(),
+            randomFrom(Completion.SUPPORTED_TASK_TYPES),
+            randomPrompt(),
+            randomAttribute()
+        );
     }
 
     @Override
@@ -29,14 +38,16 @@ protected CompletionExec mutateInstance(CompletionExec instance) throws IOExcept
         Expression inferenceId = instance.inferenceId();
         Expression prompt = instance.prompt();
         Attribute targetField = instance.targetField();
+        TaskType taskType = instance.taskType();
 
-        switch (between(0, 3)) {
+        switch (between(0, 4)) {
             case 0 -> child = randomValueOtherThan(child, () -> randomChild(0));
             case 1 -> inferenceId = randomValueOtherThan(inferenceId, this::randomInferenceId);
-            case 2 -> prompt = randomValueOtherThan(prompt, this::randomPrompt);
-            case 3 -> targetField = randomValueOtherThan(targetField, this::randomAttribute);
+            case 2 -> taskType = randomValueOtherThan(taskType, () -> randomFrom(Completion.SUPPORTED_TASK_TYPES));
+            case 3 -> prompt = randomValueOtherThan(prompt, this::randomPrompt);
+            case 4 -> targetField = randomValueOtherThan(targetField, this::randomAttribute);
         }
-        return new CompletionExec(instance.source(), child, inferenceId, prompt, targetField);
+        return new CompletionExec(instance.source(), child, inferenceId, taskType, prompt, targetField);
     }
 
     private Literal randomInferenceId() {