Skip to content

Commit aa20074

Browse files
davidkyledavidkyle
committed
[ML] Fix request format for Cohere V2 completions (elastic#131091)
(cherry picked from commit 88d765d) # Conflicts: # x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/CohereServiceUpgradeIT.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/request/v1/CohereV1CompletionRequest.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/request/v2/CohereV2CompletionRequest.java
1 parent 92547ee commit aa20074

File tree

8 files changed

+231
-8
lines changed

8 files changed

+231
-8
lines changed

‎x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/CohereServiceUpgradeIT.java‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,12 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
3838

3939
private static final String COHERE_EMBEDDINGS_ADDED = "8.13.0";
4040
private static final String COHERE_RERANK_ADDED = "8.14.0";
41+
private static final String COHERE_COMPLETIONS_ADDED_TEST_FEATURE = "gte_v8.15.0";
4142
private static final String COHERE_V2_API_ADDED_TEST_FEATURE = "inference.cohere.v2";
4243

4344
private static MockWebServer cohereEmbeddingsServer;
4445
private static MockWebServer cohereRerankServer;
46+
private static MockWebServer cohereCompletionsServer;
4547

4648
private enum ApiVersion {
4749
V1,
@@ -59,12 +61,16 @@ public static void startWebServer() throws IOException {
5961

6062
cohereRerankServer = new MockWebServer();
6163
cohereRerankServer.start();
64+
65+
cohereCompletionsServer = new MockWebServer();
66+
cohereCompletionsServer.start();
6267
}
6368

6469
@AfterClass
6570
public static void shutdown() {
6671
cohereEmbeddingsServer.close();
6772
cohereRerankServer.close();
73+
cohereCompletionsServer.close();
6874
}
6975

7076
@SuppressWarnings("unchecked")
@@ -357,6 +363,88 @@ private void assertRerank(String inferenceId) throws IOException {
357363
assertThat(inferenceMap.entrySet(), not(empty()));
358364
}
359365

366+
@SuppressWarnings("unchecked")
367+
public void testCohereCompletions() throws IOException {
368+
var completionsSupported = oldClusterHasFeature(COHERE_COMPLETIONS_ADDED_TEST_FEATURE);
369+
assumeTrue("Cohere completions not supported", completionsSupported);
370+
371+
ApiVersion oldClusterApiVersion = oldClusterHasFeature(COHERE_V2_API_ADDED_TEST_FEATURE) ? ApiVersion.V2 : ApiVersion.V1;
372+
373+
final String oldClusterId = "old-cluster-completions";
374+
375+
if (isOldCluster()) {
376+
// queue a response as PUT will call the service
377+
cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(oldClusterApiVersion)));
378+
put(oldClusterId, completionsConfig(getUrl(cohereCompletionsServer)), TaskType.COMPLETION);
379+
380+
var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
381+
assertThat(configs, hasSize(1));
382+
assertEquals("cohere", configs.get(0).get("service"));
383+
var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
384+
assertThat(serviceSettings, hasEntry("model_id", "command"));
385+
} else if (isMixedCluster()) {
386+
var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
387+
assertThat(configs, hasSize(1));
388+
assertEquals("cohere", configs.get(0).get("service"));
389+
var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
390+
assertThat(serviceSettings, hasEntry("model_id", "command"));
391+
} else if (isUpgradedCluster()) {
392+
// check old cluster model
393+
var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
394+
var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
395+
assertThat(serviceSettings, hasEntry("model_id", "command"));
396+
397+
final String newClusterId = "new-cluster-completions";
398+
{
399+
cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(oldClusterApiVersion)));
400+
var inferenceMap = inference(oldClusterId, TaskType.COMPLETION, "some text");
401+
assertThat(inferenceMap.entrySet(), not(empty()));
402+
assertVersionInPath(
403+
cohereCompletionsServer.requests().get(cohereCompletionsServer.requests().size() - 1),
404+
"chat",
405+
oldClusterApiVersion
406+
);
407+
}
408+
{
409+
// new cluster uses the V2 API
410+
cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(ApiVersion.V2)));
411+
put(newClusterId, completionsConfig(getUrl(cohereCompletionsServer)), TaskType.COMPLETION);
412+
413+
cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(ApiVersion.V2)));
414+
var inferenceMap = inference(newClusterId, TaskType.COMPLETION, "some text");
415+
assertThat(inferenceMap.entrySet(), not(empty()));
416+
assertVersionInPath(
417+
cohereCompletionsServer.requests().get(cohereCompletionsServer.requests().size() - 1),
418+
"chat",
419+
ApiVersion.V2
420+
);
421+
}
422+
423+
{
424+
// new endpoints use the V2 API which require the model to be set
425+
final String upgradedClusterNoModel = "upgraded-cluster-missing-model-id";
426+
var jsonBody = Strings.format("""
427+
{
428+
"service": "cohere",
429+
"service_settings": {
430+
"url": "%s",
431+
"api_key": "XXXX"
432+
}
433+
}
434+
""", getUrl(cohereEmbeddingsServer));
435+
436+
var e = expectThrows(ResponseException.class, () -> put(upgradedClusterNoModel, jsonBody, TaskType.COMPLETION));
437+
assertThat(
438+
e.getMessage(),
439+
containsString("Validation Failed: 1: The [service_settings.model_id] field is required for the Cohere V2 API.")
440+
);
441+
}
442+
443+
delete(oldClusterId);
444+
delete(newClusterId);
445+
}
446+
}
447+
360448
private String embeddingConfigByte(String url) {
361449
return embeddingConfigTemplate(url, "byte");
362450
}
@@ -482,4 +570,86 @@ private String rerankResponse() {
482570
""";
483571
}
484572

573+
private String completionsConfig(String url) {
574+
return Strings.format("""
575+
{
576+
"service": "cohere",
577+
"service_settings": {
578+
"api_key": "XXXX",
579+
"model_id": "command",
580+
"url": "%s"
581+
}
582+
}
583+
""", url);
584+
}
585+
586+
private String completionsResponse(ApiVersion version) {
587+
return switch (version) {
588+
case V1 -> v1CompletionsResponse();
589+
case V2 -> v2CompletionsResponse();
590+
};
591+
}
592+
593+
private String v1CompletionsResponse() {
594+
return """
595+
{
596+
"response_id": "some id",
597+
"text": "result",
598+
"generation_id": "some id",
599+
"chat_history": [
600+
{
601+
"role": "USER",
602+
"message": "some input"
603+
},
604+
{
605+
"role": "CHATBOT",
606+
"message": "v1 response from the llm"
607+
}
608+
],
609+
"finish_reason": "COMPLETE",
610+
"meta": {
611+
"api_version": {
612+
"version": "1"
613+
},
614+
"billed_units": {
615+
"input_tokens": 4,
616+
"output_tokens": 191
617+
},
618+
"tokens": {
619+
"input_tokens": 70,
620+
"output_tokens": 191
621+
}
622+
}
623+
}
624+
""";
625+
}
626+
627+
private String v2CompletionsResponse() {
628+
return """
629+
{
630+
"id": "c14c80c3-18eb-4519-9460-6c92edd8cfb4",
631+
"finish_reason": "COMPLETE",
632+
"message": {
633+
"role": "assistant",
634+
"content": [
635+
{
636+
"type": "text",
637+
"text": "v2 response from the LLM"
638+
}
639+
]
640+
},
641+
"usage": {
642+
"billed_units": {
643+
"input_tokens": 1,
644+
"output_tokens": 2
645+
},
646+
"tokens": {
647+
"input_tokens": 3,
648+
"output_tokens": 4
649+
}
650+
}
651+
}
652+
""";
653+
}
654+
485655
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/request/CohereUtils.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,16 @@ public class CohereUtils {
2828
public static final String DOCUMENTS_FIELD = "documents";
2929
public static final String EMBEDDING_TYPES_FIELD = "embedding_types";
3030
public static final String INPUT_TYPE_FIELD = "input_type";
31-
public static final String MESSAGE_FIELD = "message";
31+
public static final String V1_MESSAGE_FIELD = "message";
32+
public static final String V2_MESSAGES_FIELD = "messages";
3233
public static final String MODEL_FIELD = "model";
3334
public static final String QUERY_FIELD = "query";
35+
public static final String V2_ROLE_FIELD = "role";
3436
public static final String SEARCH_DOCUMENT = "search_document";
3537
public static final String SEARCH_QUERY = "search_query";
36-
public static final String TEXTS_FIELD = "texts";
3738
public static final String STREAM_FIELD = "stream";
39+
public static final String TEXTS_FIELD = "texts";
40+
public static final String USER_FIELD = "user";
3841

3942
public static Header createRequestSourceHeader() {
4043
return new BasicHeader(REQUEST_SOURCE_HEADER, ELASTIC_REQUEST_SOURCE);

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/request/v1/CohereV1CompletionRequest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ public CohereV1CompletionRequest(List<String> input, CohereCompletionModel model
3030
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
3131
builder.startObject();
3232
// we only allow one input for completion, so always get the first one
33-
builder.field(CohereUtils.MESSAGE_FIELD, input.get(0));
33+
builder.field(CohereUtils.V1_MESSAGE_FIELD, input.get(0));
3434
if (getModelId() != null) {
3535
builder.field(CohereUtils.MODEL_FIELD, getModelId());
3636
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/request/v2/CohereV2CompletionRequest.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,13 @@ public CohereV2CompletionRequest(List<String> input, CohereCompletionModel model
2929
@Override
3030
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
3131
builder.startObject();
32+
builder.startArray(CohereUtils.V2_MESSAGES_FIELD);
33+
builder.startObject();
34+
builder.field(CohereUtils.V2_ROLE_FIELD, CohereUtils.USER_FIELD);
3235
// we only allow one input for completion, so always get the first one
33-
builder.field(CohereUtils.MESSAGE_FIELD, input.get(0));
36+
builder.field("content", input.get(0));
37+
builder.endObject();
38+
builder.endArray();
3439
builder.field(CohereUtils.MODEL_FIELD, getModelId());
3540
builder.field(CohereUtils.STREAM_FIELD, isStreaming());
3641
builder.endObject();

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/action/CohereActionCreatorTests.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,10 @@ public void testCreate_CohereCompletionModel_WithModelSpecified() throws IOExcep
209209
assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), is("Bearer secret"));
210210

211211
var requestMap = entityAsMap(webServer.requests().get(0).getBody());
212-
assertThat(requestMap, is(Map.of("message", "abc", "model", "model", "stream", false)));
212+
assertThat(
213+
requestMap,
214+
is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "model", "stream", false))
215+
);
213216
}
214217
}
215218
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/action/CohereCompletionActionTests.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,10 @@ public void testExecute_ReturnsSuccessfulResponse_WithModelSpecified() throws IO
132132
);
133133

134134
var requestMap = entityAsMap(webServer.requests().get(0).getBody());
135-
assertThat(requestMap, is(Map.of("message", "abc", "model", "model", "stream", false)));
135+
assertThat(
136+
requestMap,
137+
is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "model", "stream", false))
138+
);
136139
}
137140
}
138141

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/request/v2/CohereV2CompletionRequestTests.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,10 @@ public void testCreateRequest() throws IOException {
4646
assertThat(httpPost.getLastHeader(CohereUtils.REQUEST_SOURCE_HEADER).getValue(), is(CohereUtils.ELASTIC_REQUEST_SOURCE));
4747

4848
var requestMap = entityAsMap(httpPost.getEntity().getContent());
49-
assertThat(requestMap, is(Map.of("message", "abc", "model", "required model id", "stream", false)));
49+
assertThat(
50+
requestMap,
51+
is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "required model id", "stream", false))
52+
);
5053
}
5154

5255
public void testDefaultUrl() {
@@ -88,6 +91,6 @@ public void testXContents() throws IOException {
8891
String xContentResult = Strings.toString(builder);
8992

9093
assertThat(xContentResult, CoreMatchers.is("""
91-
{"message":"some input","model":"model","stream":false}"""));
94+
{"messages":[{"role":"user","content":"some input"}],"model":"model","stream":false}"""));
9295
}
9396
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/response/CohereCompletionResponseEntityTests.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,42 @@ public void testFromResponse_CreatesResponseEntityForText() throws IOException {
6464
assertThat(chatCompletionResults.getResults().get(0).content(), is("result"));
6565
}
6666

67+
public void testFromResponseV2() throws IOException {
68+
String responseJson = """
69+
{
70+
"id": "abc123",
71+
"finish_reason": "COMPLETE",
72+
"message": {
73+
"role": "assistant",
74+
"content": [
75+
{
76+
"type": "text",
77+
"text": "Response from the llm"
78+
}
79+
]
80+
},
81+
"usage": {
82+
"billed_units": {
83+
"input_tokens": 1,
84+
"output_tokens": 4
85+
},
86+
"tokens": {
87+
"input_tokens": 2,
88+
"output_tokens": 5
89+
}
90+
}
91+
}
92+
""";
93+
94+
ChatCompletionResults chatCompletionResults = CohereCompletionResponseEntity.fromResponse(
95+
mock(Request.class),
96+
new HttpResult(mock(HttpResponse.class), responseJson.getBytes(StandardCharsets.UTF_8))
97+
);
98+
99+
assertThat(chatCompletionResults.getResults().size(), is(1));
100+
assertThat(chatCompletionResults.getResults().get(0).content(), is("Response from the llm"));
101+
}
102+
67103
public void testFromResponse_FailsWhenTextIsNotPresent() {
68104
String responseJson = """
69105
{

0 commit comments

Comments
 (0)