v2 messages format

Author: davidkyle

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

@@ -28,13 +28,16 @@ public class CohereUtils {
     public static final String DOCUMENTS_FIELD = "documents";
     public static final String EMBEDDING_TYPES_FIELD = "embedding_types";
     public static final String INPUT_TYPE_FIELD = "input_type";
-    public static final String MESSAGE_FIELD = "message";
+    public static final String V1_MESSAGE_FIELD = "message";
+    public static final String V2_MESSAGES_FIELD = "messages";
     public static final String MODEL_FIELD = "model";
     public static final String QUERY_FIELD = "query";
+    public static final String V2_ROLE_FIELD = "role";
     public static final String SEARCH_DOCUMENT = "search_document";
     public static final String SEARCH_QUERY = "search_query";
-    public static final String TEXTS_FIELD = "texts";
     public static final String STREAM_FIELD = "stream";
+    public static final String TEXTS_FIELD = "texts";
+    public static final String USER_FIELD = "user";
 
     public static Header createRequestSourceHeader() {
         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

@@ -30,7 +30,7 @@ public CohereV1CompletionRequest(List<String> input, CohereCompletionModel model
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         // we only allow one input for completion, so always get the first one
-        builder.field(CohereUtils.MESSAGE_FIELD, input.getFirst());
+        builder.field(CohereUtils.V1_MESSAGE_FIELD, input.getFirst());
         if (getModelId() != null) {
             builder.field(CohereUtils.MODEL_FIELD, getModelId());
         }

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

@@ -29,8 +29,13 @@ public CohereV2CompletionRequest(List<String> input, CohereCompletionModel model
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
+        builder.startArray(CohereUtils.V2_MESSAGES_FIELD);
+        builder.startObject();
+        builder.field(CohereUtils.V2_ROLE_FIELD, CohereUtils.USER_FIELD);
         // we only allow one input for completion, so always get the first one
-        builder.field(CohereUtils.MESSAGE_FIELD, input.getFirst());
+        builder.field("content", input.getFirst());
+        builder.endObject();
+        builder.endArray();
         builder.field(CohereUtils.MODEL_FIELD, getModelId());
         builder.field(CohereUtils.STREAM_FIELD, isStreaming());
         builder.endObject();

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

@@ -209,7 +209,10 @@ public void testCreate_CohereCompletionModel_WithModelSpecified() throws IOExcep
             assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), is("Bearer secret"));
 
             var requestMap = entityAsMap(webServer.requests().get(0).getBody());
-            assertThat(requestMap, is(Map.of("message", "abc", "model", "model", "stream", false)));
+            assertThat(
+                requestMap,
+                is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "model", "stream", false))
+            );
         }
     }
 }

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

@@ -132,7 +132,10 @@ public void testExecute_ReturnsSuccessfulResponse_WithModelSpecified() throws IO
             );
 
             var requestMap = entityAsMap(webServer.requests().get(0).getBody());
-            assertThat(requestMap, is(Map.of("message", "abc", "model", "model", "stream", false)));
+            assertThat(
+                requestMap,
+                is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "model", "stream", false))
+            );
         }
     }
 

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

@@ -46,7 +46,10 @@ public void testCreateRequest() throws IOException {
         assertThat(httpPost.getLastHeader(CohereUtils.REQUEST_SOURCE_HEADER).getValue(), is(CohereUtils.ELASTIC_REQUEST_SOURCE));
 
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
-        assertThat(requestMap, is(Map.of("message", "abc", "model", "required model id", "stream", false)));
+        assertThat(
+            requestMap,
+            is(Map.of("messages", List.of(Map.of("role", "user", "content", "abc")), "model", "required model id", "stream", false))
+        );
     }
 
     public void testDefaultUrl() {
@@ -88,6 +91,6 @@ public void testXContents() throws IOException {
         String xContentResult = Strings.toString(builder);
 
         assertThat(xContentResult, CoreMatchers.is("""
-            {"message":"some input","model":"model","stream":false}"""));
+            {"messages":[{"role":"user","content":"some input"}],"model":"model","stream":false}"""));
     }
 }

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

@@ -64,6 +64,42 @@ public void testFromResponse_CreatesResponseEntityForText() throws IOException {
         assertThat(chatCompletionResults.getResults().get(0).content(), is("result"));
     }
 
+    public void testFromResponseV2() throws IOException {
+        String responseJson = """
+            {
+              "id": "abc123",
+              "finish_reason": "COMPLETE",
+              "message": {
+                "role": "assistant",
+                "content": [
+                  {
+                    "type": "text",
+                    "text": "Response from the llm"
+                  }
+                ]
+              },
+              "usage": {
+                "billed_units": {
+                  "input_tokens": 1,
+                  "output_tokens": 4
+                },
+                "tokens": {
+                  "input_tokens": 2,
+                  "output_tokens": 5
+                }
+              }
+            }
+            """;
+
+        ChatCompletionResults chatCompletionResults = CohereCompletionResponseEntity.fromResponse(
+            mock(Request.class),
+            new HttpResult(mock(HttpResponse.class), responseJson.getBytes(StandardCharsets.UTF_8))
+        );
+
+        assertThat(chatCompletionResults.getResults().size(), is(1));
+        assertThat(chatCompletionResults.getResults().get(0).content(), is("Response from the llm"));
+    }
+
     public void testFromResponse_FailsWhenTextIsNotPresent() {
         String responseJson = """
             {