Skip to content

Commit 127eeb7

Browse files
author
Max Hniebergall
committed
Merge branch 'ml-inference-unified-api-elastic' of github.com:elastic/elasticsearch into ml-inference-unified-api-elastic
2 parents 6831aaa + 99d202f commit 127eeb7

File tree

5 files changed

+157
-74
lines changed

5 files changed

+157
-74
lines changed

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

Lines changed: 64 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ public record UnifiedCompletionRequest(
3333
List<Message> messages,
3434
@Nullable String model,
3535
@Nullable Long maxCompletionTokens,
36-
@Nullable Stop stop,
36+
@Nullable List<String> stop,
3737
@Nullable Float temperature,
3838
@Nullable ToolChoice toolChoice,
3939
@Nullable List<Tool> tools,
@@ -49,7 +49,7 @@ public sealed interface Content extends NamedWriteable permits ContentObjects, C
4949
(List<Message>) args[0],
5050
(String) args[1],
5151
(Long) args[2],
52-
(Stop) args[3],
52+
(List<String>) args[3],
5353
(Float) args[4],
5454
(ToolChoice) args[5],
5555
(List<Tool>) args[6],
@@ -61,7 +61,9 @@ public sealed interface Content extends NamedWriteable permits ContentObjects, C
6161
PARSER.declareObjectArray(constructorArg(), Message.PARSER::apply, new ParseField("messages"));
6262
PARSER.declareString(optionalConstructorArg(), new ParseField("model"));
6363
PARSER.declareLong(optionalConstructorArg(), new ParseField("max_completion_tokens"));
64-
PARSER.declareField(optionalConstructorArg(), (p, c) -> parseStop(p), new ParseField("stop"), ObjectParser.ValueType.VALUE_ARRAY);
64+
// PARSER.declareField(optionalConstructorArg(), (p, c) -> parseStop(p), new ParseField("stop"),
65+
// ObjectParser.ValueType.VALUE_ARRAY);
66+
PARSER.declareStringArray(optionalConstructorArg(), new ParseField("stop"));
6567
PARSER.declareFloat(optionalConstructorArg(), new ParseField("temperature"));
6668
PARSER.declareField(
6769
optionalConstructorArg(),
@@ -78,9 +80,9 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
7880
new NamedWriteableRegistry.Entry(Content.class, ContentObjects.NAME, ContentObjects::new),
7981
new NamedWriteableRegistry.Entry(Content.class, ContentString.NAME, ContentString::new),
8082
new NamedWriteableRegistry.Entry(ToolChoice.class, ToolChoiceObject.NAME, ToolChoiceObject::new),
81-
new NamedWriteableRegistry.Entry(ToolChoice.class, ToolChoiceString.NAME, ToolChoiceString::new),
82-
new NamedWriteableRegistry.Entry(Stop.class, StopValues.NAME, StopValues::new),
83-
new NamedWriteableRegistry.Entry(Stop.class, StopString.NAME, StopString::new)
83+
new NamedWriteableRegistry.Entry(ToolChoice.class, ToolChoiceString.NAME, ToolChoiceString::new)
84+
// new NamedWriteableRegistry.Entry(Stop.class, StopValues.NAME, StopValues::new),
85+
// new NamedWriteableRegistry.Entry(Stop.class, StopString.NAME, StopString::new)
8486
);
8587
}
8688

@@ -93,7 +95,7 @@ public UnifiedCompletionRequest(StreamInput in) throws IOException {
9395
in.readCollectionAsImmutableList(Message::new),
9496
in.readOptionalString(),
9597
in.readOptionalVLong(),
96-
in.readOptionalNamedWriteable(Stop.class),
98+
in.readOptionalStringCollectionAsList(),
9799
in.readOptionalFloat(),
98100
in.readOptionalNamedWriteable(ToolChoice.class),
99101
in.readOptionalCollectionAsList(Tool::new),
@@ -106,7 +108,7 @@ public void writeTo(StreamOutput out) throws IOException {
106108
out.writeCollection(messages);
107109
out.writeOptionalString(model);
108110
out.writeOptionalVLong(maxCompletionTokens);
109-
out.writeOptionalNamedWriteable(stop);
111+
out.writeOptionalStringCollection(stop);
110112
out.writeOptionalFloat(temperature);
111113
out.writeOptionalNamedWriteable(toolChoice);
112114
out.writeOptionalCollection(tools);
@@ -279,60 +281,60 @@ public void writeTo(StreamOutput out) throws IOException {
279281
}
280282
}
281283

282-
private static Stop parseStop(XContentParser parser) throws IOException {
283-
var token = parser.currentToken();
284-
if (token == XContentParser.Token.START_ARRAY) {
285-
var parsedStopValues = XContentParserUtils.parseList(parser, XContentParser::text);
286-
return new StopValues(parsedStopValues);
287-
} else if (token == XContentParser.Token.VALUE_STRING) {
288-
return StopString.of(parser);
289-
}
290-
291-
throw new XContentParseException("Unsupported token [" + token + "]");
292-
}
293-
294-
public sealed interface Stop extends NamedWriteable permits StopString, StopValues {}
295-
296-
public record StopString(String value) implements Stop, NamedWriteable {
297-
public static final String NAME = "stop_string";
298-
299-
public static StopString of(XContentParser parser) throws IOException {
300-
var content = parser.text();
301-
return new StopString(content);
302-
}
303-
304-
public StopString(StreamInput in) throws IOException {
305-
this(in.readString());
306-
}
307-
308-
@Override
309-
public void writeTo(StreamOutput out) throws IOException {
310-
out.writeString(value);
311-
}
312-
313-
@Override
314-
public String getWriteableName() {
315-
return NAME;
316-
}
317-
}
318-
319-
public record StopValues(List<String> values) implements Stop, NamedWriteable {
320-
public static final String NAME = "stop_values";
321-
322-
public StopValues(StreamInput in) throws IOException {
323-
this(in.readStringCollectionAsImmutableList());
324-
}
325-
326-
@Override
327-
public void writeTo(StreamOutput out) throws IOException {
328-
out.writeStringCollection(values);
329-
}
330-
331-
@Override
332-
public String getWriteableName() {
333-
return NAME;
334-
}
335-
}
284+
// private static Stop parseStop(XContentParser parser) throws IOException {
285+
// var token = parser.currentToken();
286+
// if (token == XContentParser.Token.START_ARRAY) {
287+
// var parsedStopValues = XContentParserUtils.parseList(parser, XContentParser::text);
288+
// return new StopValues(parsedStopValues);
289+
// } else if (token == XContentParser.Token.VALUE_STRING) {
290+
// return StopString.of(parser);
291+
// }
292+
//
293+
// throw new XContentParseException("Unsupported token [" + token + "]");
294+
// }
295+
296+
// public sealed interface Stop extends NamedWriteable permits StopString, StopValues {}
297+
//
298+
// public record StopString(String value) implements Stop, NamedWriteable {
299+
// public static final String NAME = "stop_string";
300+
//
301+
// public static StopString of(XContentParser parser) throws IOException {
302+
// var content = parser.text();
303+
// return new StopString(content);
304+
// }
305+
//
306+
// public StopString(StreamInput in) throws IOException {
307+
// this(in.readString());
308+
// }
309+
//
310+
// @Override
311+
// public void writeTo(StreamOutput out) throws IOException {
312+
// out.writeString(value);
313+
// }
314+
//
315+
// @Override
316+
// public String getWriteableName() {
317+
// return NAME;
318+
// }
319+
// }
320+
//
321+
// public record StopValues(List<String> values) implements Stop, NamedWriteable {
322+
// public static final String NAME = "stop_values";
323+
//
324+
// public StopValues(StreamInput in) throws IOException {
325+
// this(in.readStringCollectionAsImmutableList());
326+
// }
327+
//
328+
// @Override
329+
// public void writeTo(StreamOutput out) throws IOException {
330+
// out.writeStringCollection(values);
331+
// }
332+
//
333+
// @Override
334+
// public String getWriteableName() {
335+
// return NAME;
336+
// }
337+
// }
336338

337339
private static ToolChoice parseToolChoice(XContentParser parser) throws IOException {
338340
var token = parser.currentToken();

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/InferenceActionRequestTests.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,6 @@ protected InferenceAction.Request createTestInstance() {
4141
return new InferenceAction.Request(
4242
randomFrom(TaskType.values()),
4343
randomAlphaOfLength(6),
44-
// null,
4544
randomAlphaOfLengthOrNull(10),
4645
randomList(1, 5, () -> randomAlphaOfLength(8)),
4746
randomMap(0, 3, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionActionRequestTests.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.core.inference.action;
9+
10+
import org.elasticsearch.TransportVersion;
11+
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
12+
import org.elasticsearch.common.io.stream.Writeable;
13+
import org.elasticsearch.core.TimeValue;
14+
import org.elasticsearch.inference.TaskType;
15+
import org.elasticsearch.inference.UnifiedCompletionRequest;
16+
import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
17+
18+
import java.io.IOException;
19+
import java.util.List;
20+
21+
import static org.hamcrest.Matchers.is;
22+
23+
public class UnifiedCompletionActionRequestTests extends AbstractBWCWireSerializationTestCase<UnifiedCompletionAction.Request> {
24+
25+
public void testValidation_ReturnsException_When_UnifiedCompletionRequestMessage_Is_Null() {
26+
var request = new UnifiedCompletionAction.Request(
27+
"inference_id",
28+
TaskType.COMPLETION,
29+
UnifiedCompletionRequest.of(null),
30+
TimeValue.timeValueSeconds(10)
31+
);
32+
var exception = request.validate();
33+
assertThat(exception.getMessage(), is("Validation Failed: 1: Field [messages] cannot be null;"));
34+
}
35+
36+
public void testValidation_ReturnsException_When_UnifiedCompletionRequest_Is_EmptyArray() {
37+
var request = new UnifiedCompletionAction.Request(
38+
"inference_id",
39+
TaskType.COMPLETION,
40+
UnifiedCompletionRequest.of(List.of()),
41+
TimeValue.timeValueSeconds(10)
42+
);
43+
var exception = request.validate();
44+
assertThat(exception.getMessage(), is("Validation Failed: 1: Field [messages] cannot be an empty array;"));
45+
}
46+
47+
public void testValidation_ReturnsException_When_TaskType_IsNot_Completion() {
48+
var request = new UnifiedCompletionAction.Request(
49+
"inference_id",
50+
TaskType.SPARSE_EMBEDDING,
51+
UnifiedCompletionRequest.of(List.of(UnifiedCompletionRequestTests.randomMessage())),
52+
TimeValue.timeValueSeconds(10)
53+
);
54+
var exception = request.validate();
55+
assertThat(exception.getMessage(), is("Validation Failed: 1: Field [taskType] must be [completion];"));
56+
}
57+
58+
@Override
59+
protected UnifiedCompletionAction.Request mutateInstanceForVersion(UnifiedCompletionAction.Request instance, TransportVersion version) {
60+
return instance;
61+
}
62+
63+
@Override
64+
protected Writeable.Reader<UnifiedCompletionAction.Request> instanceReader() {
65+
return UnifiedCompletionAction.Request::new;
66+
}
67+
68+
@Override
69+
protected UnifiedCompletionAction.Request createTestInstance() {
70+
return new UnifiedCompletionAction.Request(
71+
randomAlphaOfLength(10),
72+
randomFrom(TaskType.values()),
73+
UnifiedCompletionRequestTests.randomUnifiedCompletionRequest(),
74+
TimeValue.timeValueMillis(randomLongBetween(1, 2048))
75+
);
76+
}
77+
78+
@Override
79+
protected UnifiedCompletionAction.Request mutateInstance(UnifiedCompletionAction.Request instance) throws IOException {
80+
return randomValueOtherThan(instance, this::createTestInstance);
81+
}
82+
83+
@Override
84+
protected NamedWriteableRegistry getNamedWriteableRegistry() {
85+
return new NamedWriteableRegistry(UnifiedCompletionRequest.getNamedWriteables());
86+
}
87+
}

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionRequestTests.java‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public void testParseAllFields() throws IOException {
9696
),
9797
"gpt-4o",
9898
100L,
99-
new UnifiedCompletionRequest.StopValues(List.of("stop")),
99+
List.of("stop"),
100100
0.1F,
101101
new UnifiedCompletionRequest.ToolChoiceObject(
102102
"function",
@@ -161,7 +161,7 @@ public void testParsing() throws IOException {
161161
),
162162
"gpt-4o",
163163
null,
164-
new UnifiedCompletionRequest.StopString("none"),
164+
List.of("none"),
165165
null,
166166
new UnifiedCompletionRequest.ToolChoiceString("auto"),
167167
List.of(
@@ -227,14 +227,12 @@ public static UnifiedCompletionRequest.ToolCall.FunctionField randomToolCallFunc
227227
return new UnifiedCompletionRequest.ToolCall.FunctionField(randomAlphaOfLength(10), randomAlphaOfLength(10));
228228
}
229229

230-
public static UnifiedCompletionRequest.Stop randomStopOrNull() {
230+
public static List<String> randomStopOrNull() {
231231
return randomBoolean() ? randomStop() : null;
232232
}
233233

234-
public static UnifiedCompletionRequest.Stop randomStop() {
235-
return randomBoolean()
236-
? new UnifiedCompletionRequest.StopString(randomAlphaOfLength(10))
237-
: new UnifiedCompletionRequest.StopValues(randomList(5, () -> randomAlphaOfLength(10)));
234+
public static List<String> randomStop() {
235+
return randomList(5, () -> randomAlphaOfLength(10));
238236
}
239237

240238
public static UnifiedCompletionRequest.ToolChoice randomToolChoiceOrNull() {

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,10 +115,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
115115
builder.field(NUMBER_OF_RETURNED_CHOICES_FIELD, 1);
116116

117117
if (unifiedRequest.stop() != null) {
118-
switch (unifiedRequest.stop()) {
119-
case UnifiedCompletionRequest.StopString stopString -> builder.field(STOP_FIELD, stopString.value());
120-
case UnifiedCompletionRequest.StopValues stopValues -> builder.field(STOP_FIELD, stopValues.values());
121-
}
118+
builder.field(STOP_FIELD, unifiedRequest.stop());
122119
}
123120
if (unifiedRequest.temperature() != null) {
124121
builder.field(TEMPERATURE_FIELD, unifiedRequest.temperature());

0 commit comments

Comments
 (0)