Skip to content

Commit 3dfb8f5

Browse files
author
Max Hniebergall
committed
tests for StreamingUnifiedChatCompletionResultsTests toXContentChunked
1 parent 0fb9a17 commit 3dfb8f5

File tree

2 files changed

+265
-4
lines changed

2 files changed

+265
-4
lines changed

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/StreamingUnifiedChatCompletionResults.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,9 +145,9 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
145145
if (usage != null) {
146146
usageIterator = Iterators.concat(
147147
ChunkedToXContentHelper.startObject(USAGE_FIELD),
148-
ChunkedToXContentHelper.field("completion_tokens", usage.completionTokens()),
149-
ChunkedToXContentHelper.field("prompt_tokens", usage.promptTokens()),
150-
ChunkedToXContentHelper.field("total_tokens", usage.totalTokens()),
148+
ChunkedToXContentHelper.field(COMPLETION_TOKENS_FIELD, usage.completionTokens()),
149+
ChunkedToXContentHelper.field(PROMPT_TOKENS_FIELD, usage.promptTokens()),
150+
ChunkedToXContentHelper.field(TOTAL_TOKENS_FIELD, usage.totalTokens()),
151151
ChunkedToXContentHelper.endObject()
152152
);
153153
}

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/results/StreamingUnifiedChatCompletionResultsTests.java‎

Lines changed: 262 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,273 @@
33
* or more contributor license agreements. Licensed under the Elastic License
44
* 2.0; you may not use this file except in compliance with the Elastic License
55
* 2.0.
6+
*
7+
* this file was contributed to by a generative AI
68
*/
79

810
package org.elasticsearch.xpack.core.inference.results;
911

12+
import org.elasticsearch.common.Strings;
1013
import org.elasticsearch.test.ESTestCase;
14+
import org.elasticsearch.xcontent.XContentBuilder;
15+
import org.elasticsearch.xcontent.json.JsonXContent;
16+
17+
import java.io.IOException;
18+
import java.util.ArrayDeque;
19+
import java.util.Deque;
20+
import java.util.List;
1121

1222
public class StreamingUnifiedChatCompletionResultsTests extends ESTestCase {
13-
// See OpenAiUnifiedStreamingProcessorTests.java
23+
24+
public void testResultstoXContentChunked() throws IOException {
25+
String expected = """
26+
{
27+
"id": "chunk1",
28+
"choices": [
29+
{
30+
"delta": {
31+
"content": "example_content",
32+
"refusal": "example_refusal",
33+
"role": "assistant",
34+
"tool_calls": [
35+
{
36+
"index": 1,
37+
"id": "tool1",
38+
"function": {
39+
"arguments": "example_arguments",
40+
"name": "example_function"
41+
},
42+
"type": "function"
43+
}
44+
]
45+
},
46+
"finish_reason": "example_reason",
47+
"index": 0
48+
}
49+
],
50+
"model": "example_model",
51+
"object": "example_object",
52+
"usage": {
53+
"completion_tokens": 10,
54+
"prompt_tokens": 5,
55+
"total_tokens": 15
56+
}
57+
}
58+
""";
59+
60+
StreamingUnifiedChatCompletionResults.ChatCompletionChunk chunk = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk(
61+
"chunk1",
62+
List.of(
63+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice(
64+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta(
65+
"example_content",
66+
"example_refusal",
67+
"assistant",
68+
List.of(
69+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(
70+
1,
71+
"tool1",
72+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(
73+
"example_arguments",
74+
"example_function"
75+
),
76+
"function"
77+
)
78+
)
79+
),
80+
"example_reason",
81+
0
82+
)
83+
),
84+
"example_model",
85+
"example_object",
86+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Usage(10, 5, 15)
87+
);
88+
89+
Deque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> deque = new ArrayDeque<>();
90+
deque.add(chunk);
91+
StreamingUnifiedChatCompletionResults.Results results = new StreamingUnifiedChatCompletionResults.Results(deque);
92+
XContentBuilder builder = JsonXContent.contentBuilder();
93+
results.toXContentChunked(null).forEachRemaining(xContent -> {
94+
try {
95+
xContent.toXContent(builder, null);
96+
} catch (IOException e) {
97+
throw new RuntimeException(e);
98+
}
99+
});
100+
101+
assertEquals(expected.replaceAll("\\s+", ""), Strings.toString(builder.prettyPrint()).trim());
102+
}
103+
104+
public void testChatCompletionChunkToXContentChunked() throws IOException {
105+
String expected = """
106+
{
107+
"id": "chunk1",
108+
"choices": [
109+
{
110+
"delta": {
111+
"content": "example_content",
112+
"refusal": "example_refusal",
113+
"role": "assistant",
114+
"tool_calls": [
115+
{
116+
"index": 1,
117+
"id": "tool1",
118+
"function": {
119+
"arguments": "example_arguments",
120+
"name": "example_function"
121+
},
122+
"type": "function"
123+
}
124+
]
125+
},
126+
"finish_reason": "example_reason",
127+
"index": 0
128+
}
129+
],
130+
"model": "example_model",
131+
"object": "example_object",
132+
"usage": {
133+
"completion_tokens": 10,
134+
"prompt_tokens": 5,
135+
"total_tokens": 15
136+
}
137+
}
138+
""";
139+
140+
StreamingUnifiedChatCompletionResults.ChatCompletionChunk chunk = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk(
141+
"chunk1",
142+
List.of(
143+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice(
144+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta(
145+
"example_content",
146+
"example_refusal",
147+
"assistant",
148+
List.of(
149+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(
150+
1,
151+
"tool1",
152+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(
153+
"example_arguments",
154+
"example_function"
155+
),
156+
"function"
157+
)
158+
)
159+
),
160+
"example_reason",
161+
0
162+
)
163+
),
164+
"example_model",
165+
"example_object",
166+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Usage(10, 5, 15)
167+
);
168+
169+
XContentBuilder builder = JsonXContent.contentBuilder();
170+
chunk.toXContentChunked(null).forEachRemaining(xContent -> {
171+
try {
172+
xContent.toXContent(builder, null);
173+
} catch (IOException e) {
174+
throw new RuntimeException(e);
175+
}
176+
});
177+
178+
assertEquals(expected.replaceAll("\\s+", ""), Strings.toString(builder.prettyPrint()).trim());
179+
}
180+
181+
public void testChoiceToXContentChunked() throws IOException {
182+
String expected = """
183+
{
184+
"delta": {
185+
"content": "example_content",
186+
"refusal": "example_refusal",
187+
"role": "assistant",
188+
"tool_calls": [
189+
{
190+
"index": 1,
191+
"id": "tool1",
192+
"function": {
193+
"arguments": "example_arguments",
194+
"name": "example_function"
195+
},
196+
"type": "function"
197+
}
198+
]
199+
},
200+
"finish_reason": "example_reason",
201+
"index": 0
202+
}
203+
""";
204+
205+
StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice choice =
206+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice(
207+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta(
208+
"example_content",
209+
"example_refusal",
210+
"assistant",
211+
List.of(
212+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(
213+
1,
214+
"tool1",
215+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(
216+
"example_arguments",
217+
"example_function"
218+
),
219+
"function"
220+
)
221+
)
222+
),
223+
"example_reason",
224+
0
225+
);
226+
227+
XContentBuilder builder = JsonXContent.contentBuilder();
228+
choice.toXContentChunked(null).forEachRemaining(xContent -> {
229+
try {
230+
xContent.toXContent(builder, null);
231+
} catch (IOException e) {
232+
throw new RuntimeException(e);
233+
}
234+
});
235+
236+
assertEquals(expected.replaceAll("\\s+", ""), Strings.toString(builder.prettyPrint()).trim());
237+
}
238+
239+
public void testToolCallToXContentChunked() throws IOException {
240+
String expected = """
241+
{
242+
"index": 1,
243+
"id": "tool1",
244+
"function": {
245+
"arguments": "example_arguments",
246+
"name": "example_function"
247+
},
248+
"type": "function"
249+
}
250+
""";
251+
252+
StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall toolCall =
253+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(
254+
1,
255+
"tool1",
256+
new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(
257+
"example_arguments",
258+
"example_function"
259+
),
260+
"function"
261+
);
262+
263+
XContentBuilder builder = JsonXContent.contentBuilder();
264+
toolCall.toXContentChunked(null).forEachRemaining(xContent -> {
265+
try {
266+
xContent.toXContent(builder, null);
267+
} catch (IOException e) {
268+
throw new RuntimeException(e);
269+
}
270+
});
271+
272+
assertEquals(expected.replaceAll("\\s+", ""), Strings.toString(builder.prettyPrint()).trim());
273+
}
274+
14275
}

0 commit comments

Comments
 (0)