Browse Source

[Inference API] Refactor OpenAiChatCompletionAction only 1 document input (#108551)

Tim Grein 1 year ago
parent
commit
ebbdc7a0ce

+ 7 - 6
x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiChatCompletionAction.java

@@ -44,16 +44,17 @@ public class OpenAiChatCompletionAction implements ExecutableAction {
 
     @Override
     public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
-        if (inferenceInputs instanceof DocumentsOnlyInput docsOnlyInput) {
-            if (docsOnlyInput.getInputs().size() > 1) {
-                listener.onFailure(new ElasticsearchStatusException("OpenAI completions only accepts 1 input", RestStatus.BAD_REQUEST));
-                return;
-            }
-        } else {
+        if (inferenceInputs instanceof DocumentsOnlyInput == false) {
             listener.onFailure(new ElasticsearchStatusException("Invalid inference input type", RestStatus.INTERNAL_SERVER_ERROR));
             return;
         }
 
+        var docsOnlyInput = (DocumentsOnlyInput) inferenceInputs;
+        if (docsOnlyInput.getInputs().size() > 1) {
+            listener.onFailure(new ElasticsearchStatusException("OpenAI completions only accepts 1 input", RestStatus.BAD_REQUEST));
+            return;
+        }
+
         try {
             ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(errorMessage, listener);