Make API_COMPLETIONS_PATH public and add unit tests for Ai21ChatCompletionRequestEntity and Ai21ChatCompletionRequest

Author: Jan-Kazlouski-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ai21/completion/Ai21ChatCompletionModel.java

@@ -28,7 +28,7 @@
  * This class extends RateLimitGroupingModel to handle rate limiting based on model and API key.
  */
 public class Ai21ChatCompletionModel extends Ai21Model {
-    private static final String API_COMPLETIONS_PATH = "https://api.ai21.com/studio/v1/chat/completions";
+    public static final String API_COMPLETIONS_PATH = "https://api.ai21.com/studio/v1/chat/completions";
 
     /**
      * Constructor for Ai21ChatCompletionModel.

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ai21/request/Ai21ChatCompletionRequestEntityTests.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services.ai21.request;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.inference.UnifiedCompletionRequest;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.json.JsonXContent;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+import org.elasticsearch.xpack.inference.services.ai21.completion.Ai21ChatCompletionModel;
+
+import java.io.IOException;
+import java.util.ArrayList;
+
+import static org.elasticsearch.xpack.inference.services.ai21.completion.Ai21ChatCompletionModelTests.createCompletionModel;
+
+public class Ai21ChatCompletionRequestEntityTests extends ESTestCase {
+    private static final String ROLE = "user";
+
+    public void testModelUserFieldsSerialization() throws IOException {
+        UnifiedCompletionRequest.Message message = new UnifiedCompletionRequest.Message(
+            new UnifiedCompletionRequest.ContentString("Hello, world!"),
+            ROLE,
+            null,
+            null
+        );
+        var messageList = new ArrayList<UnifiedCompletionRequest.Message>();
+        messageList.add(message);
+
+        var unifiedRequest = UnifiedCompletionRequest.of(messageList);
+
+        UnifiedChatInput unifiedChatInput = new UnifiedChatInput(unifiedRequest, true);
+        Ai21ChatCompletionModel model = createCompletionModel("api-key", "test-model");
+
+        Ai21ChatCompletionRequestEntity entity = new Ai21ChatCompletionRequestEntity(unifiedChatInput, model);
+
+        XContentBuilder builder = JsonXContent.contentBuilder();
+        entity.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        String expectedJson = """
+            {
+                "messages": [{
+                        "content": "Hello, world!",
+                        "role": "user"
+                    }
+                ],
+                "model": "test-model",
+                "n": 1,
+                "stream": true,
+                "stream_options": {
+                    "include_usage": true
+                }
+            }
+            """;
+        assertEquals(XContentHelper.stripWhitespace(expectedJson), Strings.toString(builder));
+    }
+
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ai21/request/Ai21ChatCompletionRequestTests.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services.ai21.request;
+
+import org.apache.http.client.methods.HttpPost;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+import org.elasticsearch.xpack.inference.services.ai21.completion.Ai21ChatCompletionModel;
+import org.elasticsearch.xpack.inference.services.ai21.completion.Ai21ChatCompletionModelTests;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap;
+import static org.hamcrest.Matchers.aMapWithSize;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+
+public class Ai21ChatCompletionRequestTests extends ESTestCase {
+
+    public void testCreateRequest_WithStreaming() throws IOException {
+        var request = createRequest("secret", randomAlphaOfLength(15), "model", true);
+        var httpRequest = request.createHttpRequest();
+
+        assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
+        var httpPost = (HttpPost) httpRequest.httpRequestBase();
+
+        var requestMap = entityAsMap(httpPost.getEntity().getContent());
+        assertThat(requestMap.get("stream"), is(true));
+    }
+
+    public void testTruncate_DoesNotReduceInputTextSize() throws IOException {
+        String input = randomAlphaOfLength(5);
+        var request = createRequest("secret", input, "model", true);
+        var truncatedRequest = request.truncate();
+        assertThat(request.getURI().toString(), is(Ai21ChatCompletionModel.API_COMPLETIONS_PATH));
+
+        var httpRequest = truncatedRequest.createHttpRequest();
+        assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
+
+        var httpPost = (HttpPost) httpRequest.httpRequestBase();
+        var requestMap = entityAsMap(httpPost.getEntity().getContent());
+        assertThat(requestMap, aMapWithSize(5));
+
+        // We do not truncate for AI21 chat completions
+        assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", input))));
+        assertThat(requestMap.get("model"), is("model"));
+        assertThat(requestMap.get("n"), is(1));
+        assertTrue((Boolean) requestMap.get("stream"));
+        assertThat(requestMap.get("stream_options"), is(Map.of("include_usage", true)));
+    }
+
+    public void testTruncationInfo_ReturnsNull() {
+        var request = createRequest("secret", randomAlphaOfLength(5), "model", true);
+        assertNull(request.getTruncationInfo());
+    }
+
+    public static Ai21ChatCompletionRequest createRequest(String apiKey, String input, @Nullable String model) {
+        return createRequest(apiKey, input, model, false);
+    }
+
+    public static Ai21ChatCompletionRequest createRequest(String apiKey, String input, @Nullable String model, boolean stream) {
+        var chatCompletionModel = Ai21ChatCompletionModelTests.createCompletionModel(apiKey, model);
+        return new Ai21ChatCompletionRequest(new UnifiedChatInput(List.of(input), "user", stream), chatCompletionModel);
+    }
+
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/request/completion/MistralChatCompletionRequestTests.java

@@ -49,7 +49,7 @@ public void testTruncate_DoesNotReduceInputTextSize() throws IOException {
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
         assertThat(requestMap, aMapWithSize(4));
 
-        // We do not truncate for Hugging Face chat completions
+        // We do not truncate for Mistral chat completions
         assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", input))));
         assertThat(requestMap.get("model"), is("model"));
         assertThat(requestMap.get("n"), is(1));