Skip to content

Refactor response handlers to improve error handling and streamline mid-stream error processing #128923

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Refactor response handlers to improve error handling and streamline mid-stream error processing #128923

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4371bbd
Refactor response handlers to improve error handling and streamline m…
Jan-Kazlouski-elastic Jun 3, 2025
99d4216
Merge branch 'main' into feature/streaming-error-handling-refactoring
Jan-Kazlouski-elastic Jun 6, 2025
9b696be
Merge branch 'main' into feature/streaming-error-handling-refactoring
Jan-Kazlouski-elastic Jun 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
231 changes: 225 additions & 6 deletions .../main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,12 @@
import org.elasticsearch.common.Strings;
import org.elasticsearch.inference.InferenceServiceResults;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
import org.elasticsearch.xpack.inference.external.http.HttpResult;
import org.elasticsearch.xpack.inference.external.request.Request;
import org.elasticsearch.xpack.inference.logging.ThrottlerManager;

import java.util.Locale;
import java.util.Objects;
import java.util.function.Function;

Expand All @@ -34,17 +36,23 @@ public abstract class BaseResponseHandler implements ResponseHandler {
public static final String SERVER_ERROR_OBJECT = "Received an error response";
public static final String BAD_REQUEST = "Received a bad request status code";
public static final String METHOD_NOT_ALLOWED = "Received a method not allowed status code";
protected static final String ERROR_TYPE = "error";
protected static final String STREAM_ERROR = "stream_error";

protected final String requestType;
protected final ResponseParser parseFunction;
private final Function<HttpResult, ErrorResponse> errorParseFunction;
private final boolean canHandleStreamingResponses;

public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorResponse> errorParseFunction) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This constructor is not used anywhere but children classes, making it protected for that reason.

protected BaseResponseHandler(
String requestType,
ResponseParser parseFunction,
Function<HttpResult, ErrorResponse> errorParseFunction
) {
this(requestType, parseFunction, errorParseFunction, false);
}

public BaseResponseHandler(
protected BaseResponseHandler(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This constructor is not used anywhere but children classes, making it protected for that reason.

String requestType,
ResponseParser parseFunction,
Function<HttpResult, ErrorResponse> errorParseFunction,
Expand Down Expand Up @@ -109,19 +117,230 @@ private void checkForErrorObject(Request request, HttpResult result) {
}

protected Exception buildError(String message, Request request, HttpResult result) {
var errorEntityMsg = errorParseFunction.apply(result);
return buildError(message, request, result, errorEntityMsg);
var errorResponse = errorParseFunction.apply(result);
return buildError(message, request, result, errorResponse);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

errorParseFunction creates ErrorResponse instance, not a message. errorEntityMsg variable naming doesn't really make sense here. Fixed.

}

protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
var responseStatusCode = result.response().getStatusLine().getStatusCode();
return new ElasticsearchStatusException(
errorMessage(message, request, result, errorResponse, responseStatusCode),
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed unused parameter

errorMessage(message, request, errorResponse, responseStatusCode),
toRestStatus(responseStatusCode)
);
}

protected String errorMessage(String message, Request request, HttpResult result, ErrorResponse errorResponse, int statusCode) {
/**
* Builds an error for a streaming request with a custom error type.
* This method is used when an error response is received from the external service.
* Only streaming requests support this format, and it should be used when the error response.
*
* @param message the error message to include in the exception
* @param request the request that caused the error
* @param result the HTTP result containing the error response
* @param errorResponse the parsed error response from the HTTP result
* @param errorResponseClass the class of the expected error response type
* @return an instance of {@link UnifiedChatCompletionException} with details from the error response
*/
protected UnifiedChatCompletionException buildChatCompletionError(
String message,
Request request,
HttpResult result,
ErrorResponse errorResponse,
Class<? extends ErrorResponse> errorResponseClass
) {
assert request.isStreaming() : "Only streaming requests support this format";
var statusCode = result.response().getStatusLine().getStatusCode();
var errorMessage = errorMessage(message, request, errorResponse, statusCode);
var restStatus = toRestStatus(statusCode);

return buildChatCompletionError(errorResponse, errorMessage, restStatus, errorResponseClass);
}

/**
* Builds a {@link UnifiedChatCompletionException} for a streaming request.
* This method is used when an error response is received from the external service.
* Only streaming requests should use this method.
*
* @param errorResponse the error response parsed from the HTTP result
* @param errorMessage the error message to include in the exception
* @param restStatus the REST status code of the response
* @param errorResponseClass the class of the expected error response type
* @return an instance of {@link UnifiedChatCompletionException} with details from the error response
*/
protected UnifiedChatCompletionException buildChatCompletionError(
ErrorResponse errorResponse,
String errorMessage,
RestStatus restStatus,
Class<? extends ErrorResponse> errorResponseClass
) {
if (errorResponseClass.isInstance(errorResponse)) {
return buildProviderSpecificChatCompletionError(errorResponse, errorMessage, restStatus);
} else {
return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
}
}

/**
* Builds a custom {@link UnifiedChatCompletionException} for a streaming request.
* This method is called when a specific error response is found in the HTTP result.
* It must be implemented by subclasses to handle specific error response formats.
* Only streaming requests should use this method.
*
* @param errorResponse the error response parsed from the HTTP result
* @param errorMessage the error message to include in the exception
* @param restStatus the REST status code of the response
* @return an instance of {@link UnifiedChatCompletionException} with details from the error response
*/
protected UnifiedChatCompletionException buildProviderSpecificChatCompletionError(
ErrorResponse errorResponse,
String errorMessage,
RestStatus restStatus
) {
throw new UnsupportedOperationException(
"Custom error handling is not implemented. Please override buildProviderSpecificChatCompletionError method."
);
}

/**
* Builds a default {@link UnifiedChatCompletionException} for a streaming request.
* This method is used when an error response is received but no specific error handling is implemented.
* Only streaming requests should use this method.
*
* @param errorResponse the error response parsed from the HTTP result
* @param errorMessage the error message to include in the exception
* @param restStatus the REST status code of the response
* @return an instance of {@link UnifiedChatCompletionException} with details from the error response
*/
protected UnifiedChatCompletionException buildDefaultChatCompletionError(
ErrorResponse errorResponse,
String errorMessage,
RestStatus restStatus
) {
return new UnifiedChatCompletionException(
restStatus,
errorMessage,
createErrorType(errorResponse),
restStatus.name().toLowerCase(Locale.ROOT)
);
}

/**
* Builds a mid-stream error for a streaming request.
* This method is used when an error occurs while processing a streaming response.
* It must be implemented by subclasses to handle specific error response formats.
* Only streaming requests should use this method.
*
* @param inferenceEntityId the ID of the inference entity
* @param message the error message
* @param e the exception that caused the error, can be null
* @return a {@link UnifiedChatCompletionException} representing the mid-stream error
*/
public UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need this here? It seems like we always override this and call the buildMidStreamChatCompletionError implementation below that takes an exception class.

Can we remove this from the base class and have the methods be private/public within the implementing classes?

throw new UnsupportedOperationException(
"Mid-stream error handling is not implemented. Please override buildMidStreamChatCompletionError method."
);
}

/**
* Builds a mid-stream error for a streaming request with a custom error type.
* This method is used when an error occurs while processing a streaming response and allows for custom error handling.
* Only streaming requests should use this method.
*
* @param inferenceEntityId the ID of the inference entity
* @param message the error message
* @param e the exception that caused the error, can be null
* @param errorResponseClass the class of the expected error response type
* @return a {@link UnifiedChatCompletionException} representing the mid-stream error
*/
protected UnifiedChatCompletionException buildMidStreamChatCompletionError(
String inferenceEntityId,
String message,
Exception e,
Class<? extends ErrorResponse> errorResponseClass
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of passing in a class here how about we either:

  • Take a function that accepts a string and returns some error class
  • Take a builder-like interface that has a method that accepts a string and returns an error class

) {
// Extract the error response from the message using the provided method
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me know if commenting here is excessive.

var errorResponse = extractMidStreamChatCompletionErrorResponse(message);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this class is getting kind of confusing with the various protected methods. I think there are a few places where we can move protected methods into parameters potentially which might make things clearer because we won't need a base class implementation that throws an error.

After the parameter change I mentioned above we can call the method/function we passed in.

// Check if the error response matches the expected type
if (errorResponseClass.isInstance(errorResponse)) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think what this is trying to do is figure out if we were able to parse the json string or not. If we were able to parse json string then we'll return the child error class. For the places that I looked, if we can't parse the string we'll return ErrorResponse.UNDEFINED_ERROR.

So instead of checking for the object instance I think we can use errorResponse.errorStructureFound().

// If it matches, we can build a custom mid-stream error exception
return buildProviderSpecificMidStreamChatCompletionError(inferenceEntityId, errorResponse);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we were able to parse error response we now want to convert it into UnifiedChatCompletionException. I wonder if we can move the functionality into the error classes.

How about we create a new class/interface that extends ErrorResponse which requires these methods:

    UnifiedChatCompletionException toUnifiedChatCompletionException();

Then I think we can have the streaming related errors extend/implement this class/interface instead and then we can do return parsedError.toUnifiedChatCompletionException()

} else if (e != null) {
// If the error response does not match, we can still return an exception based on the original throwable
return UnifiedChatCompletionException.fromThrowable(e);
} else {
// If no specific error response is found, we return a default mid-stream error
return buildDefaultMidStreamChatCompletionError(inferenceEntityId, errorResponse);
}
}

/**
* Builds a custom mid-stream {@link UnifiedChatCompletionException} for a streaming request.
* This method is called when a specific error response is found in the message.
* It must be implemented by subclasses to handle specific error response formats.
* Only streaming requests should use this method.
*
* @param inferenceEntityId the ID of the inference entity
* @param errorResponse the error response parsed from the message
* @return an instance of {@link UnifiedChatCompletionException} with details from the error response
*/
protected UnifiedChatCompletionException buildProviderSpecificMidStreamChatCompletionError(
String inferenceEntityId,
ErrorResponse errorResponse
) {
throw new UnsupportedOperationException(
"Mid-stream error handling is not implemented for this response handler. "
+ "Please override buildProviderSpecificMidStreamChatCompletionError method."
);
}

/**
* Builds a default mid-stream error for a streaming request.
* This method is used when no specific error response is found in the message.
* Only streaming requests should use this method.
*
* @param inferenceEntityId the ID of the inference entity
* @param errorResponse the error response extracted from the message
* @return a {@link UnifiedChatCompletionException} representing the default mid-stream error
*/
protected UnifiedChatCompletionException buildDefaultMidStreamChatCompletionError(
String inferenceEntityId,
ErrorResponse errorResponse
) {
return new UnifiedChatCompletionException(
RestStatus.INTERNAL_SERVER_ERROR,
format("%s for request from inference entity id [%s]", SERVER_ERROR_OBJECT, inferenceEntityId),
createErrorType(errorResponse),
STREAM_ERROR
);
}

/**
* Extracts the mid-stream error response from the message.
* This method is used to parse the error response from a streaming message.
* It must be implemented by subclasses to handle specific error response formats.
* Only streaming requests should use this method.
*
* @param message the message containing the error response
* @return an {@link ErrorResponse} object representing the mid-stream error
*/
protected ErrorResponse extractMidStreamChatCompletionErrorResponse(String message) {
throw new UnsupportedOperationException(
"Mid-stream error extraction is not implemented. Please override extractMidStreamChatCompletionErrorResponse method."
);
}

/**
* Creates a string representation of the error type based on the provided ErrorResponse.
* This method is used to generate a human-readable error type for logging or exception messages.
*
* @param errorResponse the ErrorResponse object
* @return a string representing the error type
*/
protected static String createErrorType(ErrorResponse errorResponse) {
return errorResponse != null ? errorResponse.getClass().getSimpleName() : "unknown";
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved this from several handlers so it can be used directly.

}

protected String errorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
return (errorResponse == null
|| errorResponse.errorStructureFound() == false
|| Strings.isNullOrEmpty(errorResponse.getErrorMessage()))
Expand Down
Loading