Skip to content

Commit 59c6738

Browse files
jonathan-buttnerjonathan-buttner
committed
Fixing tests
1 parent 60a6c48 commit 59c6738

File tree

10 files changed

+38
-149
lines changed

10 files changed

+38
-149
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/UnifiedChatCompletionErrorResponseUtils.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,13 @@
2121

2222
public class UnifiedChatCompletionErrorResponseUtils {
2323

24+
/**
25+
* Creates a {@link UnifiedChatCompletionErrorParserContract} that parses the error response as a string.
26+
* This is useful for cases where the error response is too complicated to parse.
27+
*
28+
* @param type The type of the error, used for categorization.
29+
* @return A {@link UnifiedChatCompletionErrorParserContract} instance.
30+
*/
2431
public static UnifiedChatCompletionErrorParserContract createErrorParserWithStringify(String type) {
2532
return new UnifiedChatCompletionErrorParserContract() {
2633
@Override
@@ -42,12 +49,27 @@ public UnifiedChatCompletionErrorResponse parse(String result) {
4249
};
4350
}
4451

52+
/**
53+
* Creates a {@link UnifiedChatCompletionErrorParserContract} that uses a {@link ConstructingObjectParser} to parse the error response.
54+
* This is useful for cases where the error response can be parsed into an object.
55+
*
56+
* @param objectParser The {@link ConstructingObjectParser} to use for parsing the error response.
57+
* @return A {@link UnifiedChatCompletionErrorParserContract} instance.
58+
*/
4559
public static UnifiedChatCompletionErrorParserContract createErrorParserWithObjectParser(
4660
ConstructingObjectParser<Optional<UnifiedChatCompletionErrorResponse>, Void> objectParser
4761
) {
4862
return new UnifiedChatCompletionErrorParser<>((parser) -> objectParser.apply(parser, null));
4963
}
5064

65+
/**
66+
* Creates a {@link UnifiedChatCompletionErrorParserContract} that uses a generic parser function to parse the error response.
67+
* This is useful for cases where the error response can be parsed using custom logic, typically when parsing from a map.
68+
*
69+
* @param genericParser The function that takes an {@link XContentParser} and returns an {@link Optional<UnifiedChatCompletionErrorResponse>}.
70+
* @param <E> The type of exception that the parser can throw.
71+
* @return A {@link UnifiedChatCompletionErrorParserContract} instance.
72+
*/
5173
public static <E extends Exception> UnifiedChatCompletionErrorParserContract createErrorParserWithGenericParser(
5274
CheckedFunction<XContentParser, Optional<UnifiedChatCompletionErrorResponse>, E> genericParser
5375
) {
@@ -90,4 +112,6 @@ private static <E extends Exception> UnifiedChatCompletionErrorResponse executeG
90112

91113
return UnifiedChatCompletionErrorResponse.UNDEFINED_ERROR;
92114
}
115+
116+
private UnifiedChatCompletionErrorResponseUtils() {}
93117
}

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

Lines changed: 5 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,62 +7,22 @@
77

88
package org.elasticsearch.xpack.inference.services.ai21.completion;
99

10-
import org.elasticsearch.rest.RestStatus;
11-
import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
12-
import org.elasticsearch.xpack.inference.external.http.HttpResult;
13-
import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
1410
import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
15-
import org.elasticsearch.xpack.inference.external.request.Request;
11+
import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorParserContract;
12+
import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorResponseUtils;
1613
import org.elasticsearch.xpack.inference.services.openai.OpenAiUnifiedChatCompletionResponseHandler;
1714

18-
import java.util.Locale;
19-
20-
import static org.elasticsearch.core.Strings.format;
21-
2215
/**
2316
* Handles streaming chat completion responses and error parsing for AI21 inference endpoints.
2417
* Adapts the OpenAI handler to support AI21's error schema.
2518
*/
2619
public class Ai21ChatCompletionResponseHandler extends OpenAiUnifiedChatCompletionResponseHandler {
2720

2821
private static final String AI_21_ERROR = "ai21_error";
22+
private static final UnifiedChatCompletionErrorParserContract AI_21_ERROR_PARSER = UnifiedChatCompletionErrorResponseUtils
23+
.createErrorParserWithStringify(AI_21_ERROR);
2924

3025
public Ai21ChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
31-
super(requestType, parseFunction, ErrorResponse::fromResponse);
32-
}
33-
34-
@Override
35-
protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
36-
assert request.isStreaming() : "Only streaming requests support this format";
37-
var responseStatusCode = result.response().getStatusLine().getStatusCode();
38-
var errorMessage = constructErrorMessage(message, request, errorResponse, responseStatusCode);
39-
var restStatus = toRestStatus(responseStatusCode);
40-
return new UnifiedChatCompletionException(restStatus, errorMessage, AI_21_ERROR, restStatus.name().toLowerCase(Locale.ROOT));
41-
}
42-
43-
protected Exception buildMidStreamError(Request request, String message, Exception e) {
44-
var errorResponse = ErrorResponse.fromString(message);
45-
if (errorResponse.errorStructureFound()) {
46-
return new UnifiedChatCompletionException(
47-
RestStatus.INTERNAL_SERVER_ERROR,
48-
format(
49-
"%s for request from inference entity id [%s]. Error message: [%s]",
50-
SERVER_ERROR_OBJECT,
51-
request.getInferenceEntityId(),
52-
errorResponse.getErrorMessage()
53-
),
54-
AI_21_ERROR,
55-
"stream_error"
56-
);
57-
} else if (e != null) {
58-
return UnifiedChatCompletionException.fromThrowable(e);
59-
} else {
60-
return new UnifiedChatCompletionException(
61-
RestStatus.INTERNAL_SERVER_ERROR,
62-
format("%s for request from inference entity id [%s]", SERVER_ERROR_OBJECT, request.getInferenceEntityId()),
63-
createErrorType(errorResponse),
64-
"stream_error"
65-
);
66-
}
26+
super(requestType, parseFunction, AI_21_ERROR_PARSER::parse, AI_21_ERROR_PARSER);
6727
}
6828
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaChatCompletionResponseHandler.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
1111
import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorParserContract;
1212
import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorResponseUtils;
13-
import org.elasticsearch.xpack.inference.services.llama.response.LlamaErrorResponse;
1413
import org.elasticsearch.xpack.inference.services.openai.OpenAiUnifiedChatCompletionResponseHandler;
1514

1615
/**
@@ -42,6 +41,6 @@ public class LlamaChatCompletionResponseHandler extends OpenAiUnifiedChatComplet
4241
* @param parseFunction the function to parse the response
4342
*/
4443
public LlamaChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
45-
super(requestType, parseFunction, LlamaErrorResponse::fromResponse, LLAMA_STREAM_ERROR_PARSER);
44+
super(requestType, parseFunction, LLAMA_STREAM_ERROR_PARSER::parse, LLAMA_STREAM_ERROR_PARSER);
4645
}
4746
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaCompletionResponseHandler.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77

88
package org.elasticsearch.xpack.inference.services.llama.completion;
99

10+
import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
1011
import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
11-
import org.elasticsearch.xpack.inference.services.llama.response.LlamaErrorResponse;
1212
import org.elasticsearch.xpack.inference.services.openai.OpenAiChatCompletionResponseHandler;
1313

1414
/**
@@ -24,6 +24,6 @@ public class LlamaCompletionResponseHandler extends OpenAiChatCompletionResponse
2424
* @param parseFunction The function to parse the response.
2525
*/
2626
public LlamaCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
27-
super(requestType, parseFunction, LlamaErrorResponse::fromResponse);
27+
super(requestType, parseFunction, ErrorResponse::fromResponse);
2828
}
2929
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/embeddings/LlamaEmbeddingsResponseHandler.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77

88
package org.elasticsearch.xpack.inference.services.llama.embeddings;
99

10+
import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
1011
import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
11-
import org.elasticsearch.xpack.inference.services.llama.response.LlamaErrorResponse;
1212
import org.elasticsearch.xpack.inference.services.openai.OpenAiResponseHandler;
1313

1414
/**
@@ -24,6 +24,6 @@ public class LlamaEmbeddingsResponseHandler extends OpenAiResponseHandler {
2424
* @param parseFunction the function to parse the response
2525
*/
2626
public LlamaEmbeddingsResponseHandler(String requestType, ResponseParser parseFunction) {
27-
super(requestType, parseFunction, LlamaErrorResponse::fromResponse, false);
27+
super(requestType, parseFunction, ErrorResponse::fromResponse, false);
2828
}
2929
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/response/LlamaErrorResponse.java‎

Lines changed: 0 additions & 59 deletions
This file was deleted.

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/MistralUnifiedChatCompletionResponseHandler.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,13 @@
99

1010
import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
1111
import org.elasticsearch.xpack.inference.services.mistral.response.MistralErrorResponseHelper;
12-
import org.elasticsearch.xpack.inference.services.openai.OpenAiChatCompletionResponseHandler;
12+
import org.elasticsearch.xpack.inference.services.openai.OpenAiUnifiedChatCompletionResponseHandler;
1313

1414
/**
1515
* Handles streaming chat completion responses and error parsing for Mistral inference endpoints.
1616
* Adapts the OpenAI handler to support Mistral's error schema.
1717
*/
18-
public class MistralUnifiedChatCompletionResponseHandler extends OpenAiChatCompletionResponseHandler {
18+
public class MistralUnifiedChatCompletionResponseHandler extends OpenAiUnifiedChatCompletionResponseHandler {
1919

2020
/**
2121
* Constructs a MistralUnifiedChatCompletionResponseHandler with the specified request type and response parser.
@@ -24,6 +24,6 @@ public class MistralUnifiedChatCompletionResponseHandler extends OpenAiChatCompl
2424
* @param parseFunction The function to parse the response.
2525
*/
2626
public MistralUnifiedChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
27-
super(requestType, parseFunction, MistralErrorResponseHelper::fromResponse);
27+
super(requestType, parseFunction, MistralErrorResponseHelper::fromResponse, MistralErrorResponseHelper.ERROR_PARSER);
2828
}
2929
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ai21/Ai21ServiceTests.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,6 @@ public void testMidStreamUnifiedCompletionError() throws Exception {
359359
testStreamError(XContentHelper.stripWhitespace("""
360360
{
361361
"error": {
362-
"code": "stream_error",
363362
"message": "Received an error response for request from inference entity id [id].\
364363
Error message: [{\\"error\\": {\\"message\\": \\"400: Invalid value: Model 'ai213.12:3b' not found\\"}}]",
365364
"type": "ai21_error"

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/response/LlamaErrorResponseTests.java‎

Lines changed: 0 additions & 34 deletions
This file was deleted.

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -344,7 +344,7 @@ public void testUnifiedCompletionInfer() throws Exception {
344344
}
345345
}
346346

347-
public void testUnifiedCompletionNonStreamingNotFoundError() throws Exception {
347+
public void testUnifiedCompletionStreamingNotFoundError() throws Exception {
348348
String responseJson = """
349349
{
350350
"detail": "Not Found"

0 commit comments

Comments
 (0)