devmadhuu commented on code in PR #9915: URL: https://github.com/apache/ozone/pull/9915#discussion_r3225978467
########## hadoop-ozone/recon/src/main/java/org/apache/hadoop/ozone/recon/chatbot/agent/ChatbotAgent.java: ########## @@ -0,0 +1,800 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.hadoop.ozone.recon.chatbot.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.Inject; +import com.google.inject.Singleton; +import org.apache.hadoop.hdds.conf.OzoneConfiguration; +import org.apache.hadoop.ozone.recon.chatbot.ChatbotConfigKeys; +import org.apache.hadoop.ozone.recon.chatbot.ChatbotException; +import org.apache.hadoop.ozone.recon.chatbot.llm.LLMClient; +import org.apache.hadoop.ozone.recon.chatbot.llm.LLMClient.ChatMessage; +import org.apache.hadoop.ozone.recon.chatbot.llm.LLMClient.LLMResponse; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Main chatbot agent that orchestrates the conversation flow. + * Handles tool selection (figuring out what API to call), executing those calls, + * and summarization (feeding the data back to the LLM to write a nice answer). + */ +@Singleton +public class ChatbotAgent { + + private static final Logger LOG = LoggerFactory.getLogger(ChatbotAgent.class); + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final Pattern JSON_PATTERN = Pattern.compile("\\{.*\\}", Pattern.DOTALL); + + // A specific Recon API endpoint we want to handle carefully because it can return millions of rows. + private static final String LIST_KEYS_ENDPOINT_SUFFIX = "/keys/listKeys"; + + /** + * Allowlist of Recon API path prefixes the chatbot is permitted to call. + * + * This is the primary defence against prompt injection: even if an attacker tricks + * the LLM into outputting an arbitrary endpoint, the Java layer will reject it here + * before ToolExecutor makes any network call. Only paths listed here can ever be + * executed. The check uses prefix matching so that parameterised paths like + * /api/v1/containers/unhealthy/MISSING are covered by the /api/v1/containers entry. + */ + private static final Set<String> ALLOWED_ENDPOINT_PREFIXES = + Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + "/api/v1/clusterState", + "/api/v1/datanodes", + "/api/v1/pipelines", + "/api/v1/containers", + "/api/v1/keys", + "/api/v1/volumes", + "/api/v1/buckets", + "/api/v1/task/status", + "/api/v1/metrics", + "/api/v1/utilization", + "/api/v1/namespace", + "/api/v1/om" + ))); + + // The connection to Gemini/OpenAI + private final LLMClient llmClient; + + // The hands that execute the internal API calls + private final ToolExecutor toolExecutor; + + // The Cheat Sheet of all available APIs loaded from the .md file + private final String apiSchema; + + // Prompt preamble for tool selection — loaded from classpath resource + private final String toolSelectionPreamble; + + // System prompt for the summarization LLM call — loaded from classpath resource + private final String summarizationPrompt; + + // Template for the fallback response when no endpoint matches — loaded from classpath resource + private final String fallbackPromptTemplate; + + // Max API calls we allow per question (so the LLM doesn't DOS our server) + private final int maxToolCalls; + + + private final String defaultModel; + private final int maxRecordsPerAnswer; + private final int maxPagesPerAnswer; + private final int pageSizePerCall; + private final boolean requireSafeScope; + + @Inject + public ChatbotAgent(LLMClient llmClient, + ToolExecutor toolExecutor, + OzoneConfiguration configuration) { + this.llmClient = llmClient; + this.toolExecutor = toolExecutor; + + // Read the Schema (Cheat Sheet) from the resources' folder. + this.apiSchema = loadApiSchema(); + + // Load prompt texts from classpath resources so they can be edited as plain text + // without touching Java code. If a file is missing the method returns "" and the + // prompt builder falls back to an inline default. + this.toolSelectionPreamble = loadApiGuideFromClasspath( + "chatbot/recon-tool-selection-prompt-preamble.txt"); + this.summarizationPrompt = loadApiGuideFromClasspath( + "chatbot/recon-summarization-prompt.txt"); + this.fallbackPromptTemplate = loadApiGuideFromClasspath( + "chatbot/recon-fallback-prompt-template.txt"); + + if (!toolSelectionPreamble.isEmpty()) { + LOG.info("Loaded tool-selection prompt preamble from classpath"); + } + if (!summarizationPrompt.isEmpty()) { + LOG.info("Loaded summarization prompt from classpath"); + } + if (!fallbackPromptTemplate.isEmpty()) { + LOG.info("Loaded fallback prompt template from classpath"); + } + + // Load all the safeguards and settings from ozone-site.xml + this.maxToolCalls = configuration.getInt( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_MAX_TOOL_CALLS, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_MAX_TOOL_CALLS_DEFAULT); + this.defaultModel = configuration.get( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_DEFAULT_MODEL, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_DEFAULT_MODEL_DEFAULT); + this.maxRecordsPerAnswer = configuration.getInt( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_MAX_RECORDS, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_MAX_RECORDS_DEFAULT); + this.maxPagesPerAnswer = configuration.getInt( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_MAX_PAGES, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_MAX_PAGES_DEFAULT); + this.pageSizePerCall = configuration.getInt( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_PAGE_SIZE, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_PAGE_SIZE_DEFAULT); + this.requireSafeScope = configuration.getBoolean( + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_REQUIRE_SAFE_SCOPE, + ChatbotConfigKeys.OZONE_RECON_CHATBOT_EXEC_REQUIRE_SAFE_SCOPE_DEFAULT); + + LOG.info("ChatbotAgent initialized with model={}, maxRecords={}, " + + "maxPages={}, pageSize={}, requireSafeScope={}", + defaultModel, maxRecordsPerAnswer, maxPagesPerAnswer, + pageSizePerCall, requireSafeScope); + } + + /** + * THE MAIN ENTRY POINT. Processes a user query and returns a response. + * + * <p>API keys are always resolved server-side via + * {@link org.apache.hadoop.ozone.recon.chatbot.security.CredentialHelper} — there + * is no per-request key parameter. All internal errors (LLM failures, IO errors, etc.) + * are wrapped in {@link ChatbotException} so callers have a single typed exception + * to handle.</p> + * + * @param userQuery the user's question + * @param model the LLM model to use (null uses the configured default) + * @param provider explicit provider name (optional, e.g. "gemini", "openai") + * @return the chatbot response + * @throws ChatbotException if query processing fails for any reason + */ + public String processQuery(String userQuery, String model, String provider) + throws ChatbotException { + + // Safety check + if (userQuery == null || userQuery.trim().isEmpty()) { + throw new ChatbotException("Query cannot be empty"); + } + + // Use default model if the user didn't specify one. + String effectiveModel = (model != null && !model.isEmpty()) ? model : defaultModel; + + LOG.info("Processing query with model: {}, provider: {}", effectiveModel, provider == null ? "auto" : provider); + + try { + // STEP 1: Ask the LLM what API tools it wants to use to answer the question. + ToolCall toolCall = getToolCall(userQuery, effectiveModel, provider); + + // If the LLM doesn't know what API to call... + if (toolCall == null) { + // No suitable endpoint found + LOG.info("Tool selection result: NO_SUITABLE_ENDPOINT; using fallback"); + return handleFallback(userQuery, effectiveModel, provider); + } + + // If the user asked a general question (e.g. "What is Ozone?"), the LLM answers it directly without an API call. + if (toolCall.isDocumentationQuery()) { + LOG.info("Tool selection result: DOCUMENTATION_QUERY (no Recon API call)"); + return toolCall.getAnswer(); + } + + // STEP 2: Execute the internal Recon API calls + Map<String, Object> apiResponses; + Map<String, Object> executionMetadata = new HashMap<>(); + + // Scenario A: LLM says we need to call MULTIPLE APIs to get the answer + if (toolCall.isMultipleEndpoints()) { + + if (toolCall.getToolCalls() == null || toolCall.getToolCalls().isEmpty()) { + LOG.warn("LLM returned MULTI_ENDPOINT but no tool calls"); + return handleFallback(userQuery, effectiveModel, provider); + } + LOG.info("Tool selection result: MULTI_ENDPOINT count={}", + toolCall.getToolCalls().size()); + + // Check if the LLM asked for something dangerous (like scanning the whole cluster without a limit) + String clarification = buildClarificationForToolCalls(toolCall.getToolCalls()); + if (clarification != null) { + LOG.info("Execution policy returned clarification for multi-endpoint " + + "request: {}", clarification); + return clarification; + } + for (ToolCall selected : toolCall.getToolCalls()) { + LOG.info("Selected Recon API: method={}, endpoint={}, paramKeys={}, reasoning={}", + selected.getMethod(), + selected.getEndpoint(), + selected.getParameters() == null ? "[]" : selected.getParameters().keySet(), + selected.getReasoning()); + } + + // Execute all the API calls securely + apiResponses = executeMultipleToolCalls(toolCall.getToolCalls(), executionMetadata); + + // Scenario B: LLM says we only need ONE API call + } else { + if (toolCall.getEndpoint() == null || toolCall.getEndpoint().isEmpty()) { + LOG.warn("LLM returned SINGLE_ENDPOINT with empty endpoint"); + return handleFallback(userQuery, effectiveModel, provider); + } + LOG.info("Tool selection result: SINGLE_ENDPOINT method={}, endpoint={}, paramKeys={}, reasoning={}", + toolCall.getMethod(), + toolCall.getEndpoint(), + toolCall.getParameters() == null ? "[]" : toolCall.getParameters().keySet(), + toolCall.getReasoning()); + String clarification = validateToolCallForExecution(toolCall); + if (clarification != null) { + LOG.info("Execution policy returned clarification for endpoint {}: {}", + toolCall.getEndpoint(), clarification); + return clarification; + } + // Go fetch the data using our ToolExecutor! + ToolExecutor.ToolExecutionOutcome outcome = toolExecutor.executeToolCallWithPolicy( + toolCall.getEndpoint(), + toolCall.getMethod(), + toolCall.getParameters(), + maxRecordsPerAnswer, + maxPagesPerAnswer, + pageSizePerCall); + + // Save the raw JSON data the API returned + apiResponses = new HashMap<>(); + apiResponses.put(toolCall.getEndpoint(), outcome.getResponseBody()); + executionMetadata.put(toolCall.getEndpoint(), + createExecutionMetadataMap(outcome)); + } + + // STEP 3: Send the raw JSON data BACK to the LLM to format a nice answer + LOG.info("Summarization input prepared: endpointCount={}, endpoints={}", + apiResponses.size(), apiResponses.keySet()); + return summarizeResponse(userQuery, apiResponses, executionMetadata, effectiveModel, provider); + + } catch (Exception e) { + throw new ChatbotException("Failed to process chatbot query: " + e.getMessage(), e); + } + } + + /** + * "Step 1" Helper: Talks to the LLM and asks for a JSON object telling us which API to call. + */ + private ToolCall getToolCall(String userQuery, String model, + String provider) throws LLMClient.LLMException, IOException { + + // Build the "cheat sheet" prompt (includes the recon-api-guide.md) + String systemPrompt = buildToolSelectionPrompt(); + String userPrompt = "User Query: " + userQuery; + + List<ChatMessage> messages = new ArrayList<>(); + messages.add(new ChatMessage("system", systemPrompt)); + messages.add(new ChatMessage("user", userPrompt)); + + // Tuning the LLM: Temperature 0.1 means we want it to be very strict and robotic, not creative. + Map<String, Object> parameters = new HashMap<>(); + parameters.put("temperature", 0.1); + parameters.put("max_tokens", 8192); Review Comment: Pls document the proper reasoning behind these params and put up a test execution plan results with different varying parameter values in PR description to understand the behavior in recon chatbot context. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
