Merge branch 'ml-inference-unified-api-elastic' of github.com:elastic/elasticsearch into ml-inference-unified-api-elastic

Author: Max Hniebergall

server/src/main/java/org/elasticsearch/inference/InferenceService.java

@@ -127,9 +127,7 @@ void unifiedCompletionInfer(
     );
 
     /**
-     * Chunk long text according to {@code chunkingOptions} or the
-     * model defaults if {@code chunkingOptions} contains unset
-     * values.
+     * Chunk long text.
      *
      * @param model           The model
      * @param query           Inference query, mainly for re-ranking

server/src/main/java/org/elasticsearch/inference/TaskType.java

@@ -38,6 +38,10 @@ public static TaskType fromString(String name) {
     }
 
     public static TaskType fromStringOrStatusException(String name) {
+        if (name == null) {
+            throw new ElasticsearchStatusException("Task type must not be null", RestStatus.BAD_REQUEST);
+        }
+
         try {
             TaskType taskType = TaskType.fromString(name);
             return Objects.requireNonNull(taskType);

server/src/main/java/org/elasticsearch/inference/UnifiedCompletionRequest.java

@@ -63,8 +63,6 @@ public sealed interface Content extends NamedWriteable permits ContentObjects, C
         PARSER.declareObjectArray(constructorArg(), Message.PARSER::apply, new ParseField("messages"));
         PARSER.declareString(optionalConstructorArg(), new ParseField("model"));
         PARSER.declareLong(optionalConstructorArg(), new ParseField("max_completion_tokens"));
-        // PARSER.declareField(optionalConstructorArg(), (p, c) -> parseStop(p), new ParseField("stop"),
-        // ObjectParser.ValueType.VALUE_ARRAY);
         PARSER.declareStringArray(optionalConstructorArg(), new ParseField("stop"));
         PARSER.declareFloat(optionalConstructorArg(), new ParseField("temperature"));
         PARSER.declareField(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiUnifiedChatCompletionRequestEntity.java

@@ -43,6 +43,8 @@ public class OpenAiUnifiedChatCompletionRequestEntity implements ToXContentObjec
     private static final String TOOL_FIELD = "tools";
     private static final String TEXT_FIELD = "text";
     private static final String TYPE_FIELD = "type";
+    private static final String STREAM_OPTIONS_FIELD = "stream_options";
+    private static final String INCLUDE_USAGE_FIELD = "include_usage";
 
     private final UnifiedCompletionRequest unifiedRequest;
     private final boolean stream;
@@ -169,6 +171,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         }
 
         builder.field(STREAM_FIELD, stream);
+        if (stream) {
+            builder.startObject(STREAM_OPTIONS_FIELD);
+            builder.field(INCLUDE_USAGE_FIELD, true);
+            builder.endObject();
+        }
         builder.endObject();
 
         return builder;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/TaskTypeTests.java

@@ -0,0 +1,27 @@
+/*
+ * 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;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.test.ESTestCase;
+import org.hamcrest.Matchers;
+
+public class TaskTypeTests extends ESTestCase {
+
+    public void testFromStringOrStatusException() {
+        var exception = expectThrows(ElasticsearchStatusException.class, () -> TaskType.fromStringOrStatusException(null));
+        assertThat(exception.getMessage(), Matchers.is("Task type must not be null"));
+
+        exception = expectThrows(ElasticsearchStatusException.class, () -> TaskType.fromStringOrStatusException("blah"));
+        assertThat(exception.getMessage(), Matchers.is("Unknown task_type [blah]"));
+
+        assertThat(TaskType.fromStringOrStatusException("any"), Matchers.is(TaskType.ANY));
+    }
+
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiUnifiedChatCompletionRequestTests.java

@@ -30,7 +30,7 @@
 public class OpenAiUnifiedChatCompletionRequestTests extends ESTestCase {
 
     public void testCreateRequest_WithUrlOrganizationUserDefined() throws IOException {
-        var request = createRequest("www.google.com", "org", "secret", "abc", "model", "user");
+        var request = createRequest("www.google.com", "org", "secret", "abc", "model", "user", true);
         var httpRequest = request.createHttpRequest();
 
         assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
@@ -42,16 +42,27 @@ public void testCreateRequest_WithUrlOrganizationUserDefined() throws IOExceptio
         assertThat(httpPost.getLastHeader(ORGANIZATION_HEADER).getValue(), is("org"));
 
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
-        assertThat(requestMap, aMapWithSize(5));
+        assertRequestMapWithUser(requestMap, "user");
+    }
+
+    private void assertRequestMapWithoutUser(Map<String, Object> requestMap) {
+        assertRequestMapWithUser(requestMap, null);
+    }
+
+    private void assertRequestMapWithUser(Map<String, Object> requestMap, @Nullable String user) {
+        assertThat(requestMap, aMapWithSize(user != null ? 6 : 5));
         assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", "abc"))));
         assertThat(requestMap.get("model"), is("model"));
-        assertThat(requestMap.get("user"), is("user"));
+        if (user != null) {
+            assertThat(requestMap.get("user"), is(user));
+        }
         assertThat(requestMap.get("n"), is(1));
         assertTrue((Boolean) requestMap.get("stream"));
+        assertThat(requestMap.get("stream_options"), is(Map.of("include_usage", true)));
     }
 
     public void testCreateRequest_WithDefaultUrl() throws URISyntaxException, IOException {
-        var request = createRequest(null, "org", "secret", "abc", "model", "user");
+        var request = createRequest(null, "org", "secret", "abc", "model", "user", true);
         var httpRequest = request.createHttpRequest();
 
         assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
@@ -63,16 +74,12 @@ public void testCreateRequest_WithDefaultUrl() throws URISyntaxException, IOExce
         assertThat(httpPost.getLastHeader(ORGANIZATION_HEADER).getValue(), is("org"));
 
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
-        assertThat(requestMap, aMapWithSize(5));
-        assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", "abc"))));
-        assertThat(requestMap.get("model"), is("model"));
-        assertThat(requestMap.get("user"), is("user"));
-        assertThat(requestMap.get("n"), is(1));
-        assertTrue((Boolean) requestMap.get("stream"));
+        assertRequestMapWithUser(requestMap, "user");
+
     }
 
     public void testCreateRequest_WithDefaultUrlAndWithoutUserOrganization() throws URISyntaxException, IOException {
-        var request = createRequest(null, null, "secret", "abc", "model", null);
+        var request = createRequest(null, null, "secret", "abc", "model", null, true);
         var httpRequest = request.createHttpRequest();
 
         assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
@@ -84,14 +91,10 @@ public void testCreateRequest_WithDefaultUrlAndWithoutUserOrganization() throws
         assertNull(httpPost.getLastHeader(ORGANIZATION_HEADER));
 
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
-        assertThat(requestMap, aMapWithSize(4));
-        assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", "abc"))));
-        assertThat(requestMap.get("model"), is("model"));
-        assertThat(requestMap.get("n"), is(1));
-        assertTrue((Boolean) requestMap.get("stream"));
+        assertRequestMapWithoutUser(requestMap);
     }
 
-    public void testCreateRequest_WithStreaming() throws URISyntaxException, IOException {
+    public void testCreateRequest_WithStreaming() throws IOException {
         var request = createRequest(null, null, "secret", "abc", "model", null, true);
         var httpRequest = request.createHttpRequest();
 
@@ -103,7 +106,7 @@ public void testCreateRequest_WithStreaming() throws URISyntaxException, IOExcep
     }
 
     public void testTruncate_DoesNotReduceInputTextSize() throws URISyntaxException, IOException {
-        var request = createRequest(null, null, "secret", "abcd", "model", null);
+        var request = createRequest(null, null, "secret", "abcd", "model", null, true);
         var truncatedRequest = request.truncate();
         assertThat(request.getURI().toString(), is(OpenAiUnifiedChatCompletionRequest.buildDefaultUri().toString()));
 
@@ -112,17 +115,18 @@ public void testTruncate_DoesNotReduceInputTextSize() throws URISyntaxException,
 
         var httpPost = (HttpPost) httpRequest.httpRequestBase();
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
-        assertThat(requestMap, aMapWithSize(4));
+        assertThat(requestMap, aMapWithSize(5));
 
         // We do not truncate for OpenAi chat completions
         assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", "abcd"))));
         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(null, null, "secret", "abcd", "model", null);
+        var request = createRequest(null, null, "secret", "abcd", "model", null, true);
         assertNull(request.getTruncationInfo());
     }
 
@@ -147,7 +151,7 @@ public static OpenAiUnifiedChatCompletionRequest createRequest(
         boolean stream
     ) {
         var chatCompletionModel = OpenAiChatCompletionModelTests.createChatCompletionModel(url, org, apiKey, model, user);
-        return new OpenAiUnifiedChatCompletionRequest(new UnifiedChatInput(List.of(input), "user", true), chatCompletionModel);
+        return new OpenAiUnifiedChatCompletionRequest(new UnifiedChatInput(List.of(input), "user", stream), chatCompletionModel);
     }
 
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/BaseInferenceActionTests.java

@@ -8,11 +8,14 @@
 package org.elasticsearch.xpack.inference.rest;
 
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.RestRequestTests;
 import org.elasticsearch.rest.action.RestChunkedToXContentListener;
 import org.elasticsearch.test.rest.FakeRestRequest;
 import org.elasticsearch.test.rest.RestActionTestCase;
@@ -26,6 +29,10 @@
 import java.util.Map;
 
 import static org.elasticsearch.rest.RestRequest.Method.POST;
+import static org.elasticsearch.xpack.inference.rest.BaseInferenceAction.parseParams;
+import static org.elasticsearch.xpack.inference.rest.BaseInferenceAction.parseTimeout;
+import static org.elasticsearch.xpack.inference.rest.Paths.INFERENCE_ID;
+import static org.elasticsearch.xpack.inference.rest.Paths.TASK_TYPE_OR_INFERENCE_ID;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
@@ -56,6 +63,42 @@ private static String route(String param) {
         return "_route/" + param;
     }
 
+    public void testParseParams_ExtractsInferenceIdAndTaskType() {
+        var params = parseParams(
+            RestRequestTests.contentRestRequest("{}", Map.of(INFERENCE_ID, "id", TASK_TYPE_OR_INFERENCE_ID, TaskType.COMPLETION.toString()))
+        );
+        assertThat(params, is(new BaseInferenceAction.Params("id", TaskType.COMPLETION)));
+    }
+
+    public void testParseParams_DefaultsToTaskTypeAny_WhenInferenceId_IsMissing() {
+        var params = parseParams(
+            RestRequestTests.contentRestRequest("{}", Map.of(TASK_TYPE_OR_INFERENCE_ID, TaskType.COMPLETION.toString()))
+        );
+        assertThat(params, is(new BaseInferenceAction.Params("completion", TaskType.ANY)));
+    }
+
+    public void testParseParams_ThrowsStatusException_WhenTaskTypeIsMissing() {
+        var e = expectThrows(
+            ElasticsearchStatusException.class,
+            () -> parseParams(RestRequestTests.contentRestRequest("{}", Map.of(INFERENCE_ID, "id")))
+        );
+        assertThat(e.getMessage(), is("Task type must not be null"));
+    }
+
+    public void testParseTimeout_ReturnsTimeout() {
+        var timeout = parseTimeout(
+            RestRequestTests.contentRestRequest("{}", Map.of(InferenceAction.Request.TIMEOUT.getPreferredName(), "4s"))
+        );
+
+        assertThat(timeout, is(TimeValue.timeValueSeconds(4)));
+    }
+
+    public void testParseTimeout_ReturnsDefaultTimeout() {
+        var timeout = parseTimeout(RestRequestTests.contentRestRequest("{}", Map.of()));
+
+        assertThat(timeout, is(TimeValue.timeValueSeconds(30)));
+    }
+
     public void testUsesDefaultTimeout() {
         SetOnce<Boolean> executeCalled = new SetOnce<>();
         verifyingClient.setExecuteVerifier(((actionType, actionRequest) -> {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceActionTests.java

@@ -0,0 +1,81 @@
+/*
+ * 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.rest;
+
+import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.rest.AbstractRestChannel;
+import org.elasticsearch.rest.RestChannel;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.test.rest.FakeRestRequest;
+import org.elasticsearch.test.rest.RestActionTestCase;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
+import org.junit.Before;
+
+import static org.elasticsearch.xpack.inference.rest.BaseInferenceActionTests.createResponse;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
+
+public class RestUnifiedCompletionInferenceActionTests extends RestActionTestCase {
+
+    @Before
+    public void setUpAction() {
+        controller().registerHandler(new RestUnifiedCompletionInferenceAction());
+    }
+
+    public void testStreamIsTrue() {
+        SetOnce<Boolean> executeCalled = new SetOnce<>();
+        verifyingClient.setExecuteVerifier(((actionType, actionRequest) -> {
+            assertThat(actionRequest, instanceOf(UnifiedCompletionAction.Request.class));
+
+            var request = (UnifiedCompletionAction.Request) actionRequest;
+            assertThat(request.isStreaming(), is(true));
+
+            executeCalled.set(true);
+            return createResponse();
+        }));
+
+        var requestBody = """
+            {
+              "messages": [
+                {
+                  "content": "abc",
+                  "role": "user"
+                }
+              ]
+            }
+            """;
+
+        RestRequest inferenceRequest = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST)
+            .withPath("_inference/completion/test/_unified")
+            .withContent(new BytesArray(requestBody), XContentType.JSON)
+            .build();
+
+        final SetOnce<RestResponse> responseSetOnce = new SetOnce<>();
+        dispatchRequest(inferenceRequest, new AbstractRestChannel(inferenceRequest, true) {
+            @Override
+            public void sendResponse(RestResponse response) {
+                responseSetOnce.set(response);
+            }
+        });
+
+        // the response content will be null when there is no error
+        assertNull(responseSetOnce.get().content());
+        assertThat(executeCalled.get(), equalTo(true));
+    }
+
+    private void dispatchRequest(final RestRequest request, final RestChannel channel) {
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
+        controller().dispatchRequest(request, channel, threadContext);
+    }
+}