Fix: support Claude V3 output parsing in Generative QA Processor

Signed-off-by: Anuj Soni <[email protected]>

Author: sonianuj287

search-processors/src/main/java/org/opensearch/searchpipelines/questionanswering/generative/llm/DefaultLlmImpl.java

@@ -191,8 +191,40 @@ protected ChatCompletionOutput buildChatCompletionOutput(ModelProvider provider,
                 answers = List.of(message.get(CONNECTOR_OUTPUT_MESSAGE_CONTENT));
             }
         } else if (provider == ModelProvider.BEDROCK) {
-            answerField = "completion";
-            fillAnswersOrErrors(dataAsMap, answers, errors, answerField, errorField, defaultErrorMessageField);
+            // Handle both Claude V2 and V3 response formats
+            if (dataAsMap.containsKey("completion")) {
+                // Old Claude V2 format
+                answerField = "completion";
+                fillAnswersOrErrors(dataAsMap, answers, errors, answerField, errorField, defaultErrorMessageField);
+            } else if (dataAsMap.containsKey("content")) {
+                // New Claude V3 format
+                Object contentObj = dataAsMap.get("content");
+                if (contentObj instanceof List) {
+                    List<?> contentList = (List<?>) contentObj;
+                    if (!contentList.isEmpty()) {
+                        Object first = contentList.get(0);
+                        if (first instanceof Map) {
+                            Map<?, ?> firstMap = (Map<?, ?>) first;
+                            Object text = firstMap.get("text");
+                            if (text != null) {
+                                answers.add(text.toString());
+                            } else {
+                                errors.add("Claude V3 response missing 'text' field.");
+                            }
+                        } else {
+                            errors.add("Unexpected content format in Claude V3 response.");
+                        }
+                    } else {
+                        errors.add("Empty content list in Claude V3 response.");
+                    }
+                } else {
+                    errors.add("Unexpected type for 'content' in Claude V3 response.");
+                }
+            } else {
+                // Fallback error handling
+                errors.add("Unsupported Claude response format: " + dataAsMap.keySet());
+                log.error("Unknown Bedrock/Claude response format: {}", dataAsMap);
+            }
         } else if (provider == ModelProvider.COHERE) {
             answerField = "text";
             fillAnswersOrErrors(dataAsMap, answers, errors, answerField, errorField, defaultErrorMessageField);

search-processors/src/test/java/org/opensearch/searchpipelines/questionanswering/generative/llm/DefaultLlmImplTests.java

@@ -143,6 +143,58 @@ public void onFailure(Exception e) {
         assertTrue(mlInput.getInputDataset() instanceof RemoteInferenceInputDataSet);
     }
 
+    public void testChatCompletionApiForBedrockClaudeV3() throws Exception {
+        MachineLearningInternalClient mlClient = mock(MachineLearningInternalClient.class);
+        ArgumentCaptor<MLInput> captor = ArgumentCaptor.forClass(MLInput.class);
+        DefaultLlmImpl connector = new DefaultLlmImpl("model_id", client);
+        connector.setMlClient(mlClient);
+
+        // Claude V3-style response
+        Map<String, Object> textPart = Map.of("type", "text", "text", "Hello from Claude V3");
+        Map<String, Object> dataAsMap = Map.of("content", List.of(textPart));
+
+        ModelTensor tensor = new ModelTensor("tensor", new Number[0], new long[0], MLResultDataType.STRING, null, null, dataAsMap);
+        ModelTensorOutput mlOutput = new ModelTensorOutput(List.of(new ModelTensors(List.of(tensor))));
+        ActionFuture<MLOutput> future = mock(ActionFuture.class);
+        when(future.actionGet(anyLong())).thenReturn(mlOutput);
+        when(mlClient.predict(any(), any())).thenReturn(future);
+
+        ChatCompletionInput input = new ChatCompletionInput(
+            "bedrock/model",
+            "question",
+            Collections.emptyList(),
+            Collections.emptyList(),
+            0,
+            "prompt",
+            "instructions",
+            Llm.ModelProvider.BEDROCK,
+            null,
+            null
+        );
+
+        doAnswer(invocation -> {
+            ((ActionListener<MLOutput>) invocation.getArguments()[2]).onResponse(mlOutput);
+            return null;
+        }).when(mlClient).predict(any(), any(), any());
+
+        connector.doChatCompletion(input, new ActionListener<>() {
+            @Override
+            public void onResponse(ChatCompletionOutput output) {
+                // Verify that we parsed the Claude V3 response correctly
+                assertEquals("Hello from Claude V3", output.getAnswers().get(0));
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                fail("Claude V3 test failed: " + e.getMessage());
+            }
+        });
+
+        verify(mlClient, times(1)).predict(any(), captor.capture(), any());
+        MLInput mlInput = captor.getValue();
+        assertTrue(mlInput.getInputDataset() instanceof RemoteInferenceInputDataSet);
+    }
+
     public void testChatCompletionApiForBedrock() throws Exception {
         MachineLearningInternalClient mlClient = mock(MachineLearningInternalClient.class);
         ArgumentCaptor<MLInput> captor = ArgumentCaptor.forClass(MLInput.class);