Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 25, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c80ff5b
Apply patch logic
davidkyle Jun 25, 2025
656dc0a
Fix test logic
davidkyle Jun 25, 2025
99acc6d
Merge branch 'main' into fix-patch-test
davidkyle Jun 25, 2025
57817f7
Merge branch 'main' into fix-patch-test
prwhelan Jun 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ...rc/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...ticsearch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...org/elasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
Expand Down
24 changes: 22 additions & 2 deletions ...st/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettingsTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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));
Expand Down Expand Up @@ -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;
}
}
21 changes: 19 additions & 2 deletions ...arch/xpack/inference/services/cohere/completion/CohereCompletionServiceSettingsTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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(
Expand Down Expand Up @@ -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;
}
}
17 changes: 9 additions & 8 deletions ...lasticsearch/xpack/inference/services/cohere/rerank/CohereRerankServiceSettingsTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down