fix: multi-modal or tool related object will fail the search

Signed-off-by: Sicheng Song <[email protected]>

Author: b4sjoo

common/src/main/java/org/opensearch/ml/common/utils/LlmResultPathGenerator.java

@@ -58,33 +58,37 @@ public class LlmResultPathGenerator {
      */
     public static String generate(String outputSchemaJson) throws IOException {
         if (outputSchemaJson == null || outputSchemaJson.trim().isEmpty()) {
-            log.warn("Output schema is null or empty, cannot generate llm_result_path");
+            log.warn("[LLM_RESULT_PATH_GEN] Output schema is null or empty, cannot generate llm_result_path");
             return null;
         }
 
+        log.info("[LLM_RESULT_PATH_GEN] Starting llm_result_path auto-generation (schema size: {} bytes)", outputSchemaJson.length());
+
         try {
             JsonNode schemaRoot = MAPPER.readTree(outputSchemaJson);
 
             // Navigate to dataAsMap schema node using hardcoded path; if not found, search from root
             JsonNode searchRoot = navigateToDataAsMapSchema(schemaRoot);
             if (searchRoot == null) {
-                log.debug("No dataAsMap schema found, searching from root");
+                log.info("[LLM_RESULT_PATH_GEN] No dataAsMap schema found in standard ModelTensorOutput path, searching from schema root");
                 searchRoot = schemaRoot;
+            } else {
+                log.info("[LLM_RESULT_PATH_GEN] Found dataAsMap schema at standard path, searching within dataAsMap structure");
             }
 
             // Search for LLM output field with x-llm-output marker
             String jsonPath = findLlmTextField(searchRoot, "$");
 
             if (jsonPath == null) {
-                log.warn("Could not find field with x-llm-output marker in schema");
+                log.warn("[LLM_RESULT_PATH_GEN] ❌ No field with x-llm-output marker found in schema - will use fallback path");
                 return null;
             }
 
-            log.debug("Generated llm_result_path: {}", jsonPath);
+            log.info("[LLM_RESULT_PATH_GEN] ✅ Successfully generated llm_result_path: {}", jsonPath);
             return jsonPath;
 
         } catch (Exception e) {
-            log.error("Failed to generate llm_result_path from schema", e);
+            log.error("[LLM_RESULT_PATH_GEN] ❌ Failed to generate llm_result_path from schema", e);
             throw new OpenSearchParseException("Schema parsing error: " + e.getMessage(), e);
         }
     }
@@ -104,9 +108,15 @@ public static String generate(String outputSchemaJson) throws IOException {
      */
     private static JsonNode navigateToDataAsMapSchema(JsonNode schemaRoot) {
         if (schemaRoot == null || schemaRoot.isMissingNode()) {
+            log.debug("[LLM_RESULT_PATH_GEN] Schema root is null or missing");
             return null;
         }
 
+        log
+            .debug(
+                "[LLM_RESULT_PATH_GEN] Navigating to dataAsMap using rigid path: properties.inference_results.items.properties.output.items.properties.dataAsMap"
+            );
+
         // Follow the rigid ModelTensorOutput → ModelTensors → ModelTensor structure
         JsonNode dataAsMapSchema = schemaRoot
             .path("properties")
@@ -118,7 +128,13 @@ private static JsonNode navigateToDataAsMapSchema(JsonNode schemaRoot) {
             .path("properties")
             .path("dataAsMap");
 
-        return dataAsMapSchema.isMissingNode() ? null : dataAsMapSchema;
+        if (dataAsMapSchema.isMissingNode()) {
+            log.debug("[LLM_RESULT_PATH_GEN] dataAsMap not found at standard path");
+            return null;
+        }
+
+        log.debug("[LLM_RESULT_PATH_GEN] Successfully navigated to dataAsMap schema");
+        return dataAsMapSchema;
     }
 
     /**
@@ -148,6 +164,7 @@ private static String findLlmTextFieldWithMarker(JsonNode schemaNode, String cur
         // Check if this field has the x-llm-output marker
         JsonNode marker = schemaNode.get(LLM_OUTPUT_MARKER);
         if (marker != null && marker.isBoolean() && marker.asBoolean()) {
+            log.info("[LLM_RESULT_PATH_GEN] 🎯 Found x-llm-output marker at path: {}", currentPath);
             return currentPath;
         }
 
@@ -166,6 +183,7 @@ private static String findLlmTextFieldWithMarker(JsonNode schemaNode, String cur
                     JsonNode fieldSchema = field.getValue();
 
                     String newPath = currentPath + "." + fieldName;
+                    log.debug("[LLM_RESULT_PATH_GEN] Checking field: {}", newPath);
                     String result = findLlmTextFieldWithMarker(fieldSchema, newPath);
                     if (result != null) {
                         return result;
@@ -179,6 +197,7 @@ private static String findLlmTextFieldWithMarker(JsonNode schemaNode, String cur
             JsonNode items = schemaNode.get("items");
             if (items != null) {
                 String newPath = currentPath + "[0]";
+                log.debug("[LLM_RESULT_PATH_GEN] Navigating into array: {}", newPath);
                 String result = findLlmTextFieldWithMarker(items, newPath);
                 if (result != null) {
                     return result;

common/src/main/java/org/opensearch/ml/common/utils/message/ClaudeMessageFormatter.java

@@ -0,0 +1,173 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.utils.message;
+
+import static org.opensearch.common.xcontent.json.JsonXContent.jsonXContent;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.ml.common.transport.memorycontainer.memory.MessageInput;
+import org.opensearch.ml.common.utils.StringUtils;
+
+import lombok.extern.log4j.Log4j2;
+
+/**
+ * Message formatter for Claude models (and similar models using system_prompt parameter).
+ *
+ * <p>Format characteristics:
+ * <ul>
+ *   <li>System prompt: Placed in "system_prompt" parameter</li>
+ *   <li>Messages: Array of user/assistant messages (NO system role)</li>
+ *   <li>Content: Normalized to have "type" field</li>
+ * </ul>
+ *
+ * <p>Compatible with:
+ * <ul>
+ *   <li>Claude 3.x (Bedrock, Anthropic API)</li>
+ *   <li>Claude 4.x (Sonnet, Opus)</li>
+ *   <li>Any model with system_prompt in input schema</li>
+ * </ul>
+ *
+ * <p>Example output:
+ * <pre>
+ * {
+ *   "system_prompt": "You are a helpful assistant",
+ *   "messages": "[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}]"
+ * }
+ * </pre>
+ */
+@Log4j2
+public class ClaudeMessageFormatter implements MessageFormatter {
+
+    @Override
+    public Map<String, String> formatRequest(String systemPrompt, List<MessageInput> messages, Map<String, Object> additionalConfig) {
+        Map<String, String> parameters = new HashMap<>();
+
+        // Claude-style: system_prompt as parameter
+        if (systemPrompt != null && !systemPrompt.isBlank()) {
+            parameters.put("system_prompt", systemPrompt);
+            log.debug("System prompt added as parameter");
+        }
+
+        // Build messages array with content processing
+        try {
+            String messagesJson = buildMessagesArray(messages, additionalConfig);
+            parameters.put("messages", messagesJson);
+            log.debug("Built messages array with {} messages", messages != null ? messages.size() : 0);
+        } catch (IOException e) {
+            log.error("Failed to build messages array", e);
+            throw new RuntimeException("Failed to format Claude request", e);
+        }
+
+        return parameters;
+    }
+
+    @Override
+    public List<Map<String, Object>> processContent(List<Map<String, Object>> content) {
+        if (content == null || content.isEmpty()) {
+            return content;
+        }
+
+        return content.stream().map(this::normalizeContentObject).collect(Collectors.toList());
+    }
+
+    /**
+     * Normalize a single content object to ensure it has "type" field.
+     *
+     * <p>Rules:
+     * <ul>
+     *   <li>If object has "type" field → return as-is (standard LLM format)</li>
+     *   <li>If object lacks "type" field → wrap as {"type": "text", "text": JSON_STRING}</li>
+     * </ul>
+     *
+     * @param obj Content object to normalize
+     * @return Normalized content object with "type" field
+     */
+    private Map<String, Object> normalizeContentObject(Map<String, Object> obj) {
+        if (obj == null || obj.isEmpty()) {
+            return obj;
+        }
+
+        // Already has type field → standard format
+        if (obj.containsKey("type")) {
+            return obj;
+        }
+
+        // No type field → user-defined object, wrap as text
+        Map<String, Object> wrapped = new HashMap<>();
+        wrapped.put("type", "text");
+        String jsonText = StringUtils.toJson(obj);
+        wrapped.put("text", jsonText);
+        return wrapped;
+    }
+
+    /**
+     * Build messages JSON array from MessageInput list.
+     *
+     * <p>Includes:
+     * <ul>
+     *   <li>Optional system_prompt_message from config (added first)</li>
+     *   <li>User/assistant messages with processed content</li>
+     *   <li>Optional user_prompt_message from config (added last)</li>
+     * </ul>
+     *
+     * @param messages List of messages to include
+     * @param additionalConfig Optional config with extra messages
+     * @return JSON string representing messages array
+     * @throws IOException if JSON building fails
+     */
+    private String buildMessagesArray(List<MessageInput> messages, Map<String, Object> additionalConfig) throws IOException {
+        XContentBuilder builder = jsonXContent.contentBuilder();
+        builder.startArray();
+
+        // Optional system_prompt_message from config
+        if (additionalConfig != null && additionalConfig.containsKey("system_prompt_message")) {
+            Object systemPromptMsg = additionalConfig.get("system_prompt_message");
+            if (systemPromptMsg instanceof Map) {
+                @SuppressWarnings("unchecked")
+                Map<String, Object> msgMap = (Map<String, Object>) systemPromptMsg;
+                builder.map(msgMap);
+            }
+        }
+
+        // User messages (with content processing)
+        if (messages != null) {
+            for (MessageInput message : messages) {
+                builder.startObject();
+
+                if (message.getRole() != null) {
+                    builder.field("role", message.getRole());
+                }
+
+                // Process content to ensure type fields
+                if (message.getContent() != null) {
+                    List<Map<String, Object>> processedContent = processContent(message.getContent());
+                    builder.field("content", processedContent);
+                }
+
+                builder.endObject();
+            }
+        }
+
+        // Optional user_prompt_message from config
+        if (additionalConfig != null && additionalConfig.containsKey("user_prompt_message")) {
+            Object userPromptMsg = additionalConfig.get("user_prompt_message");
+            if (userPromptMsg instanceof Map) {
+                @SuppressWarnings("unchecked")
+                Map<String, Object> msgMap = (Map<String, Object>) userPromptMsg;
+                builder.map(msgMap);
+            }
+        }
+
+        builder.endArray();
+        return builder.toString();
+    }
+}

common/src/main/java/org/opensearch/ml/common/utils/message/MessageFormatter.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.utils.message;
+
+import java.util.List;
+import java.util.Map;
+
+import org.opensearch.ml.common.transport.memorycontainer.memory.MessageInput;
+
+/**
+ * Strategy interface for formatting LLM requests based on model requirements.
+ *
+ * <p>Each formatter implementation handles:
+ * <ul>
+ *   <li>System prompt placement (parameter vs message)</li>
+ *   <li>Content object normalization (type field enforcement)</li>
+ *   <li>Message array construction</li>
+ * </ul>
+ *
+ * <p>Implementations:
+ * <ul>
+ *   <li>{@link ClaudeMessageFormatter}: Uses system_prompt parameter (Claude, Bedrock models)</li>
+ *   <li>{@link OpenAIMessageFormatter}: Injects system message in array (OpenAI, GPT models)</li>
+ * </ul>
+ *
+ * <p>Usage:
+ * <pre>
+ * MessageFormatter formatter = MessageFormatterFactory.getFormatterForModel(modelId, cache);
+ * Map&lt;String, String&gt; params = formatter.formatRequest(systemPrompt, messages, config);
+ * </pre>
+ */
+public interface MessageFormatter {
+
+    /**
+     * Format a complete LLM request with proper system prompt placement.
+     *
+     * <p>This method handles:
+     * <ul>
+     *   <li>Placing system prompt in correct location (parameter or message)</li>
+     *   <li>Building messages array with content processing</li>
+     *   <li>Incorporating additional config messages (system_prompt_message, user_prompt_message)</li>
+     * </ul>
+     *
+     * @param systemPrompt The system prompt to inject (may be null or blank)
+     * @param messages User/assistant messages to include in the request
+     * @param additionalConfig Optional configuration containing:
+     *                         <ul>
+     *                           <li>system_prompt_message: Additional system-level message</li>
+     *                           <li>user_prompt_message: Additional user message</li>
+     *                         </ul>
+     * @return Map of request parameters ready for MLInput, typically containing:
+     *         <ul>
+     *           <li>"messages": JSON array of messages</li>
+     *           <li>"system_prompt": System prompt (Claude-style only)</li>
+     *         </ul>
+     * @throws RuntimeException if message array building fails
+     */
+    Map<String, String> formatRequest(String systemPrompt, List<MessageInput> messages, Map<String, Object> additionalConfig);
+
+    /**
+     * Process message content objects to ensure LLM compatibility.
+     *
+     * <p>Content processing rules:
+     * <ul>
+     *   <li>Objects WITH "type" field → keep as-is (standard LLM format like
+     *       {"type": "text", "text": "..."} or {"type": "image_url", "image_url": {...}})</li>
+     *   <li>Objects WITHOUT "type" field → wrap as {"type": "text", "text": JSON_STRING}
+     *       where JSON_STRING is the serialized user-defined object</li>
+     * </ul>
+     *
+     * <p>This ensures that user-defined content objects are properly formatted
+     * for LLM consumption without breaking standard multimodal content.
+     *
+     * @param content List of content objects from message (may be null or empty)
+     * @return Processed content list with all objects having "type" field
+     */
+    List<Map<String, Object>> processContent(List<Map<String, Object>> content);
+}