[ML] Apply patch logic for Cohere V2 transport changes (#129993)

davidkyle · web-flow · commit 5be4100afa9a · 2025-06-26T05:16:30.000+10:00
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java
@@ -183,7 +183,7 @@ public CohereServiceSettings(StreamInput in) throws IOException {
             rateLimitSettings = DEFAULT_RATE_LIMIT_SETTINGS;
         }
         if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             this.apiVersion = in.readEnum(CohereServiceSettings.CohereApiVersion.class);
         } else {
             this.apiVersion = CohereServiceSettings.CohereApiVersion.V1;
@@ -286,7 +286,7 @@ public void writeTo(StreamOutput out) throws IOException {
             rateLimitSettings.writeTo(out);
         }
         if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             out.writeEnum(apiVersion);
         }
     }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettings.java
@@ -103,7 +103,7 @@ public CohereCompletionServiceSettings(StreamInput in) throws IOException {
         modelId = in.readOptionalString();
         rateLimitSettings = new RateLimitSettings(in);
         if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             this.apiVersion = in.readEnum(CohereServiceSettings.CohereApiVersion.class);
         } else {
             this.apiVersion = CohereServiceSettings.CohereApiVersion.V1;
@@ -156,7 +156,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(modelId);
         rateLimitSettings.writeTo(out);
         if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             out.writeEnum(apiVersion);
         }
     }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettings.java
@@ -125,7 +125,7 @@ public CohereRerankServiceSettings(StreamInput in) throws IOException {
         }
 
         if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             this.apiVersion = in.readEnum(CohereServiceSettings.CohereApiVersion.class);
         } else {
             this.apiVersion = CohereServiceSettings.CohereApiVersion.V1;
@@ -207,7 +207,7 @@ public void writeTo(StreamOutput out) throws IOException {
             rateLimitSettings.writeTo(out);
         }
         if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
-            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
+            || out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {
             out.writeEnum(apiVersion);
         }
     }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettingsTests.java
@@ -7,15 +7,17 @@
 
 package org.elasticsearch.xpack.inference.services.cohere;
 
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.SimilarityMeasure;
-import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.ServiceFields;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
@@ -30,7 +32,7 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
 
-public class CohereServiceSettingsTests extends AbstractWireSerializingTestCase<CohereServiceSettings> {
+public class CohereServiceSettingsTests extends AbstractBWCWireSerializationTestCase<CohereServiceSettings> {
 
     public static CohereServiceSettings createRandomWithNonNullUrl() {
         return createRandom(randomAlphaOfLength(15));
@@ -359,4 +361,22 @@ public static Map<String, Object> getServiceSettingsMap(@Nullable String url, @N
 
         return map;
     }
+
+    @Override
+    protected CohereServiceSettings mutateInstanceForVersion(CohereServiceSettings instance, TransportVersion version) {
+        if (version.before(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
+            && (version.isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19) == false)) {
+            return new CohereServiceSettings(
+                instance.uri(),
+                instance.similarity(),
+                instance.dimensions(),
+                instance.maxInputTokens(),
+                instance.modelId(),
+                instance.rateLimitSettings(),
+                CohereServiceSettings.CohereApiVersion.V1
+            );
+        }
+
+        return instance;
+    }
 }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettingsTests.java
@@ -7,12 +7,14 @@
 
 package org.elasticsearch.xpack.inference.services.cohere.completion;
 
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.Writeable;
-import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.ServiceFields;
 import org.elasticsearch.xpack.inference.services.cohere.CohereServiceSettings;
@@ -25,7 +27,7 @@
 
 import static org.hamcrest.Matchers.is;
 
-public class CohereCompletionServiceSettingsTests extends AbstractWireSerializingTestCase<CohereCompletionServiceSettings> {
+public class CohereCompletionServiceSettingsTests extends AbstractBWCWireSerializationTestCase<CohereCompletionServiceSettings> {
 
     public static CohereCompletionServiceSettings createRandom() {
         return new CohereCompletionServiceSettings(
@@ -110,4 +112,19 @@ protected CohereCompletionServiceSettings createTestInstance() {
     protected CohereCompletionServiceSettings mutateInstance(CohereCompletionServiceSettings instance) throws IOException {
         return randomValueOtherThan(instance, this::createTestInstance);
     }
+
+    @Override
+    protected CohereCompletionServiceSettings mutateInstanceForVersion(CohereCompletionServiceSettings instance, TransportVersion version) {
+        if (version.before(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
+            && (version.isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19) == false)) {
+            return new CohereCompletionServiceSettings(
+                instance.uri(),
+                instance.modelId(),
+                instance.rateLimitSettings(),
+                CohereServiceSettings.CohereApiVersion.V1
+            );
+        }
+
+        return instance;
+    }
 }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettingsTests.java
@@ -88,14 +88,15 @@ protected CohereRerankServiceSettings mutateInstanceForVersion(CohereRerankServi
                 CohereServiceSettings.DEFAULT_RATE_LIMIT_SETTINGS,
                 CohereServiceSettings.CohereApiVersion.V1
             );
-        } else if (version.before(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {
-            return new CohereRerankServiceSettings(
-                instance.uri(),
-                instance.modelId(),
-                instance.rateLimitSettings(),
-                CohereServiceSettings.CohereApiVersion.V1
-            );
-        }
+        } else if (version.before(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)
+            && version.isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19) == false) {
+                return new CohereRerankServiceSettings(
+                    instance.uri(),
+                    instance.modelId(),
+                    instance.rateLimitSettings(),
+                    CohereServiceSettings.CohereApiVersion.V1
+                );
+            }
         return instance;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ public CohereServiceSettings(StreamInput in) throws IOException {`
`183`	`183`	`rateLimitSettings = DEFAULT_RATE_LIMIT_SETTINGS;`
`184`	`184`	`}`
`185`	`185`	`if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)`
`186`		`- \|\| in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {`
	`186`	`+ \|\| in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {`
`187`	`187`	`this.apiVersion = in.readEnum(CohereServiceSettings.CohereApiVersion.class);`
`188`	`188`	`} else {`
`189`	`189`	`this.apiVersion = CohereServiceSettings.CohereApiVersion.V1;`
`@@ -286,7 +286,7 @@ public void writeTo(StreamOutput out) throws IOException {`
`286`	`286`	`rateLimitSettings.writeTo(out);`
`287`	`287`	`}`
`288`	`288`	`if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)`
`289`		`- \|\| out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {`
	`289`	`+ \|\| out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {`
`290`	`290`	`out.writeEnum(apiVersion);`
`291`	`291`	`}`
`292`	`292`	`}`
Original file line number	Diff line number	Diff line change
`@@ -125,7 +125,7 @@ public CohereRerankServiceSettings(StreamInput in) throws IOException {`
`125`	`125`	`}`
`126`	`126`
`127`	`127`	`if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)`
`128`		`- \|\| in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {`
	`128`	`+ \|\| in.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {`
`129`	`129`	`this.apiVersion = in.readEnum(CohereServiceSettings.CohereApiVersion.class);`
`130`	`130`	`} else {`
`131`	`131`	`this.apiVersion = CohereServiceSettings.CohereApiVersion.V1;`
`@@ -207,7 +207,7 @@ public void writeTo(StreamOutput out) throws IOException {`
`207`	`207`	`rateLimitSettings.writeTo(out);`
`208`	`208`	`}`
`209`	`209`	`if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)`
`210`		`- \|\| out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION)) {`
	`210`	`+ \|\| out.getTransportVersion().isPatchFrom(TransportVersions.ML_INFERENCE_COHERE_API_VERSION_8_19)) {`
`211`	`211`	`out.writeEnum(apiVersion);`
`212`	`212`	`}`
`213`	`213`	`}`