From ecc74d4a21e77dc666509e65a2abf6587b71e2e7 Mon Sep 17 00:00:00 2001 From: Praful Kumar Date: Thu, 23 Jul 2026 12:09:28 +0530 Subject: [PATCH] add failoverLLM Introduce FailoverLlm with error classification and provider failover. Add Azure runtime override support for host-injected config and treat AUTH/BAD_REQUEST failures as retryable for failover. Co-authored-by: Cursor --- .../com/google/adk/models/BedrockBaseLM.java | 735 ++++++++---------- .../java/com/google/adk/models/Claude.java | 17 +- .../com/google/adk/models/LlmRegistry.java | 95 ++- .../google/adk/models/azure/AzureConfig.java | 33 +- .../adk/models/azure/AzureRestTransport.java | 679 ++++++++-------- .../adk/models/failover/FailoverConfig.java | 117 +++ .../adk/models/failover/FailoverLlm.java | 318 ++++++++ .../adk/models/failover/LlmCallMetrics.java | 123 +++ .../google/adk/models/failover/LlmError.java | 166 ++++ .../models/failover/LlmErrorClassifier.java | 214 +++++ .../models/failover/LlmFailureCategory.java | 69 ++ .../adk/models/failover/LlmHttpException.java | 54 ++ 12 files changed, 1884 insertions(+), 736 deletions(-) create mode 100644 core/src/main/java/com/google/adk/models/failover/FailoverConfig.java create mode 100644 core/src/main/java/com/google/adk/models/failover/FailoverLlm.java create mode 100644 core/src/main/java/com/google/adk/models/failover/LlmCallMetrics.java create mode 100644 core/src/main/java/com/google/adk/models/failover/LlmError.java create mode 100644 core/src/main/java/com/google/adk/models/failover/LlmErrorClassifier.java create mode 100644 core/src/main/java/com/google/adk/models/failover/LlmFailureCategory.java create mode 100644 core/src/main/java/com/google/adk/models/failover/LlmHttpException.java diff --git a/core/src/main/java/com/google/adk/models/BedrockBaseLM.java b/core/src/main/java/com/google/adk/models/BedrockBaseLM.java index 52bd472f1..05729c83f 100644 --- a/core/src/main/java/com/google/adk/models/BedrockBaseLM.java +++ b/core/src/main/java/com/google/adk/models/BedrockBaseLM.java @@ -4,13 +4,13 @@ */ package com.google.adk.models; -import static com.google.adk.models.RedbusADG.callLLMChat; import static com.google.adk.models.RedbusADG.cleanForIdentifierPattern; import static com.google.common.collect.ImmutableList.toImmutableList; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.datatype.jdk8.Jdk8Module; +import com.google.adk.models.failover.LlmHttpException; import com.google.adk.tools.BaseTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; @@ -24,6 +24,7 @@ import com.google.genai.types.Part; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.io.BufferedReader; import java.io.DataOutputStream; import java.io.IOException; @@ -103,116 +104,94 @@ private static String buildConverseUrl(String baseUrl, String model) { return base + "/model/" + model + "/converse"; } - // Corrected the logger name to use OllamaBaseLM.class - private static final Logger logger = LoggerFactory.getLogger(BedrockBaseLM.class); - - private static final String CONTINUE_OUTPUT_MESSAGE = - "Continue output. DO NOT look at this line. ONLY look at the content before this line and" - + " system instruction."; - - public BedrockBaseLM(String model) { + private static JSONArray buildSystemArray(String systemText) { + JSONArray system = new JSONArray(); + if (systemText != null && !systemText.isEmpty()) { + system.put(new JSONObject().put("text", systemText)); + } + return system; + } - super(model); + private static String mapBedrockRole(String adkRole) { + if ("model".equals(adkRole) || "assistant".equals(adkRole)) { + return "assistant"; + } + return "user"; } - /** - * @param model The model ID (e.g., "openai.gpt-oss-20b-1:0") - * @param BEDROCK_EP The base Bedrock endpoint (e.g., - * "https://bedrock-runtime.us-west-2.amazonaws.com/model") - */ - public BedrockBaseLM(String model, String BEDROCK_EP) { - super(model); - this.D_URL = BEDROCK_EP; + private static String toolUseIdForName(String name) { + return "tooluse_" + cleanForIdentifierPattern(name == null ? "tool" : name); } - @Override - public Flowable generateContent(LlmRequest llmRequest, boolean stream) { - if (stream) { - return generateContentStream(llmRequest); + private static String resolveToolUseIdForFunctionResponse( + List contents, int responseMessageIndex, String toolName) { + for (int i = responseMessageIndex - 1; i >= 0; i--) { + Content c = contents.get(i); + String role = c.role().orElse("user"); + if (!"model".equals(role) && !"assistant".equals(role)) { + continue; + } + for (Part p : c.parts().orElse(ImmutableList.of())) { + if (p.functionCall().isPresent()) { + FunctionCall fc = p.functionCall().get(); + if (toolName.equals(fc.name().orElse(""))) { + return fc.id().orElse(toolUseIdForName(toolName)); + } + } + } } + return toolUseIdForName(toolName); + } - List contents = llmRequest.contents(); - // Last content must be from the user, otherwise the model won't respond. - if (contents.isEmpty() || !Iterables.getLast(contents).role().orElse("").equals("user")) { - Content userContent = Content.fromParts(Part.fromText(CONTINUE_OUTPUT_MESSAGE)); - contents = - Stream.concat(contents.stream(), Stream.of(userContent)).collect(toImmutableList()); - } + /** Converts ADK Content to a Bedrock Converse message with proper toolUse/toolResult blocks. */ + private static JSONObject contentToBedrockMessage( + Content item, List allContents, int messageIndex) { + JSONObject messageQuantum = new JSONObject(); + messageQuantum.put("role", mapBedrockRole(item.role().orElse("user"))); + JSONArray contentArray = new JSONArray(); - String systemText = ""; - Optional configOpt = llmRequest.config(); - if (configOpt.isPresent()) { - Optional systemInstructionOpt = configOpt.get().systemInstruction(); - if (systemInstructionOpt.isPresent()) { - String extractedSystemText = - systemInstructionOpt.get().parts().orElse(ImmutableList.of()).stream() - .filter(p -> p.text().isPresent()) - .map(p -> p.text().get()) - .collect(Collectors.joining("\n")); - if (!extractedSystemText.isEmpty()) { - systemText = extractedSystemText; - } + for (Part part : item.parts().orElse(ImmutableList.of())) { + if (part.functionResponse().isPresent()) { + var fr = part.functionResponse().get(); + String name = fr.name().orElse("tool"); + JSONObject toolResult = new JSONObject(); + toolResult.put( + "toolUseId", resolveToolUseIdForFunctionResponse(allContents, messageIndex, name)); + toolResult.put("status", "success"); + JSONArray resultContent = new JSONArray(); + resultContent.put( + new JSONObject().put("json", new JSONObject(fr.response().orElse(Map.of())))); + toolResult.put("content", resultContent); + contentArray.put(new JSONObject().put("toolResult", toolResult)); + } else if (part.functionCall().isPresent()) { + var fc = part.functionCall().get(); + String name = fc.name().orElse("tool"); + JSONObject toolUse = new JSONObject(); + toolUse.put("toolUseId", fc.id().orElse(toolUseIdForName(name))); + toolUse.put("name", name); + toolUse.put("input", new JSONObject(fc.args().orElse(Map.of()))); + contentArray.put(new JSONObject().put("toolUse", toolUse)); + } else if (part.text().isPresent() && !part.text().get().isEmpty()) { + contentArray.put(new JSONObject().put("text", part.text().get())); } } - /** "messages": [ { "role": "user", "content": [{"text": "Hello"}] } ], */ - JSONArray messages = new JSONArray(); - - JSONObject llmMessageJson1 = new JSONObject(); - llmMessageJson1.put("role", "assistant"); - JSONObject txtMsg = new JSONObject(); - txtMsg.put("text", systemText); - JSONArray contentArray = new JSONArray(); - contentArray.put(txtMsg); - llmMessageJson1.put("content", contentArray); - messages.put(llmMessageJson1); // Agent system prompt is always added + if (contentArray.length() == 0) { + contentArray.put(new JSONObject().put("text", item.text() == null ? "" : item.text())); + } + messageQuantum.put("content", contentArray); + return messageQuantum; + } - llmRequest.contents().stream() - .forEach( - item -> { - JSONObject messageQuantum = new JSONObject(); - messageQuantum.put( - "role", - item.role().get().equals("model") || item.role().get().equals("assistant") - ? "assistant" - : "user"); - - if (item.parts().get().get(0).functionResponse().isPresent()) { - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put( - "text", - item.parts().get().get(0).functionResponse().get().name().get() - + " responded with these values, " - + new JSONObject( - item.parts().get().get(0).functionResponse().get().response().get()) - .toString()); - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - messageQuantum.put("content", contentArray2); - } else if (item.parts().get().get(0).functionCall().isPresent()) { - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put( - "text", - item.parts().get().get(0).functionCall().get().name().get() - + " is to be called with these arguments, " - + new JSONObject( - item.parts().get().get(0).functionCall().get().args().get()) - .toString()); - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - messageQuantum.put("content", contentArray2); - } else { - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put("text", item.text()); - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - messageQuantum.put("content", contentArray2); - } - messages.put(messageQuantum); - }); + private static JSONArray buildMessagesFromContents(List contents) { + JSONArray messages = new JSONArray(); + for (int i = 0; i < contents.size(); i++) { + messages.put(contentToBedrockMessage(contents.get(i), contents, i)); + } + return messages; + } - // Tools - // Define the required pattern for the name + private static JSONArray buildToolsFromRequest(LlmRequest llmRequest) { JSONArray functions = new JSONArray(); llmRequest .tools() @@ -220,136 +199,158 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre .forEach( tooldetail -> { BaseTool baseTool = tooldetail.getValue(); - - // Get the function declaration from the base tool Optional declarationOptional = baseTool.declaration(); - - // Skip this tool if there is no function declaration if (!declarationOptional.isPresent()) { - // Log a warning or handle appropriately System.err.println( "Skipping tool '" + baseTool.name() + "' with missing declaration."); - // continue; // If inside a loop - return; // If processing a single tool outside a loop + return; } - FunctionDeclaration functionDeclaration = declarationOptional.get(); - - // Build the top-level map representing the tool JSON structure Map toolMap = new HashMap<>(); - - // Add the tool's name and description from the function declaration toolMap.put("name", cleanForIdentifierPattern(functionDeclaration.name().get())); toolMap.put( "description", - cleanForIdentifierPattern( - functionDeclaration - .description() - .orElse(""))); // Use description from declaration, handle Optional - - // Build the 'parameters' object if parameters are defined + cleanForIdentifierPattern(functionDeclaration.description().orElse(""))); Optional parametersOptional = functionDeclaration.parameters(); if (parametersOptional.isPresent()) { Schema parametersSchema = parametersOptional.get(); - Map parametersMap = new HashMap<>(); - parametersMap.put( - "type", "object"); // Function parameters schema type is typically "object" - - // Build the 'properties' map within 'parameters' + parametersMap.put("type", "object"); Optional> propertiesOptional = parametersSchema.properties(); if (propertiesOptional.isPresent()) { Map propertiesMap = new HashMap<>(); - // Create ObjectMapper instance once for the loop ObjectMapper objectMapper = new ObjectMapper(); - objectMapper.registerModule( - new Jdk8Module()); // Register module for Java 8 Optionals, etc. - + objectMapper.registerModule(new Jdk8Module()); propertiesOptional .get() .forEach( (key, schema) -> { - // Convert the library's Schema object for a parameter to a generic Map Map schemaMap = objectMapper.convertValue( schema, new TypeReference>() {}); - - // Apply your custom logic to update the type string - // !!! This function updateTypeString(schemaMap) is required and not - // provided !!! - updateTypeString( - schemaMap); // Ensure this modifies schemaMap in place or returns - // the modified map - + updateTypeString(schemaMap); propertiesMap.put(key, schemaMap); }); parametersMap.put("properties", propertiesMap); } - - // Add the 'required' list within 'parameters' if present parametersSchema .required() - .ifPresent( - requiredList -> - parametersMap.put( - "required", requiredList)); // Assuming required() returns - // Optional> - - // Add the completed 'parameters' map to the main tool map + .ifPresent(requiredList -> parametersMap.put("required", requiredList)); JSONObject inputSchema = new JSONObject(); inputSchema.put("json", parametersMap); toolMap.put("inputSchema", inputSchema); } - - // Convert the complete tool map into an org.json.JSONObject JSONObject jsonToolW = new JSONObject(); + jsonToolW.put("toolSpec", new JSONObject(toolMap)); + functions.put(jsonToolW); + }); + return functions; + } + + private static void putToolConfig(JSONObject payload, JSONArray tools) { + if (tools != null && tools.length() > 0) { + JSONObject toolConfig = new JSONObject(); + toolConfig.put("tools", tools); + payload.put("toolConfig", toolConfig); + } + } - JSONObject jsonTool = new JSONObject(toolMap); + private static JSONObject buildConversePayload( + JSONArray system, JSONArray messages, JSONArray tools) { + JSONObject payload = new JSONObject(); + if (system != null && system.length() > 0) { + payload.put("system", system); + } + payload.put("messages", messages); + putToolConfig(payload, tools); + return payload; + } - jsonToolW.put("toolSpec", jsonTool); + // Corrected the logger name to use OllamaBaseLM.class + private static final Logger logger = LoggerFactory.getLogger(BedrockBaseLM.class); - // Add the generated tool JSON object to your functions list/array - functions.put(jsonToolW); - }); + private static final String CONTINUE_OUTPUT_MESSAGE = + "Continue output. DO NOT look at this line. ONLY look at the content before this line and" + + " system instruction."; - // Check if the tool is executed, then parse and response. + public BedrockBaseLM(String model) { - logger.debug("functions: {}", functions.toString(1)); + super(model); + } + + /** + * @param model The model ID (e.g., "openai.gpt-oss-20b-1:0") + * @param BEDROCK_EP The base Bedrock endpoint (e.g., + * "https://bedrock-runtime.us-west-2.amazonaws.com/model") + */ + public BedrockBaseLM(String model, String BEDROCK_EP) { + super(model); + this.D_URL = BEDROCK_EP; + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + if (stream) { + return generateContentStream(llmRequest).subscribeOn(Schedulers.io()); + } + // Defer the blocking HTTP call until subscription so it is cancellable and bounded by any + // downstream timeout, enabling latency-based failover in FailoverLlm. + return Flowable.defer(() -> generateContentSyncInternal(llmRequest)) + .subscribeOn(Schedulers.io()); + } + + private Flowable generateContentSyncInternal(LlmRequest llmRequest) { + List contents = llmRequest.contents(); + // Last content must be from the user, otherwise the model won't respond. + if (contents.isEmpty() || !Iterables.getLast(contents).role().orElse("").equals("user")) { + Content userContent = Content.fromParts(Part.fromText(CONTINUE_OUTPUT_MESSAGE)); + contents = + Stream.concat(contents.stream(), Stream.of(userContent)).collect(toImmutableList()); + } + + String systemText = ""; + Optional configOpt = llmRequest.config(); + if (configOpt.isPresent()) { + Optional systemInstructionOpt = configOpt.get().systemInstruction(); + if (systemInstructionOpt.isPresent()) { + String extractedSystemText = + systemInstructionOpt.get().parts().orElse(ImmutableList.of()).stream() + .filter(p -> p.text().isPresent()) + .map(p -> p.text().get()) + .collect(Collectors.joining("\n")); + if (!extractedSystemText.isEmpty()) { + systemText = extractedSystemText; + } + } + } - String modelId = - this.model(); // "devstral";//"llama3.2:3b-instruct-q2_K";//"llama3.2"; // The 1b doesn't - // support tool + JSONArray system = buildSystemArray(systemText); + JSONArray messages = buildMessagesFromContents(contents); - // If last user response has the function reponse, then function calla is not needed. - boolean LAST_RESP_TOOl_EXECUTED = - Iterables.getLast(Iterables.getLast(contents).parts().get()).functionResponse().isPresent(); + JSONArray functions = buildToolsFromRequest(llmRequest); + logger.debug("functions: {}", functions.toString(1)); + String modelId = this.model(); + + // toolConfig is required whenever messages contain toolUse/toolResult blocks (incl. post-tool + // turns). JSONObject agentresponse = - callLLMChat( - modelId, - messages, - LAST_RESP_TOOl_EXECUTED - ? null - : (functions.length() > 0 - ? functions - : null)); // Tools/functions can not be of 0 length - - // Attempt usage metadata extraction & logging + callLLMChat(modelId, system, messages, functions.length() > 0 ? functions : null); + + GenerateContentResponseUsageMetadata usageMetadata = null; try { - GenerateContentResponseUsageMetadata usage = getUsageMetadata(agentresponse); - if (usage != null) { + usageMetadata = getUsageMetadata(agentresponse); + if (usageMetadata != null) { logger.info( "Non-streaming token counts: prompt={}, completion={}, total={}", - usage.promptTokenCount().orElse(0), - usage.candidatesTokenCount().orElse(0), - usage.totalTokenCount().orElse(0)); + usageMetadata.promptTokenCount().orElse(0), + usageMetadata.candidatesTokenCount().orElse(0), + usageMetadata.totalTokenCount().orElse(0)); } } catch (Exception e) { logger.debug("Usage metadata parsing failed (non-critical)", e); } - // Bedrock Converse API: output.message, or message/Message at top level. - // Message (capital M) can be a String in error responses (e.g. "Authentication failed"). JSONObject responseQuantum = extractMessageObject(agentresponse); if (responseQuantum == null) { String detail = "Response keys: " + agentresponse.keySet(); @@ -373,35 +374,36 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre + detail); } - // Check if tool call is required - // Tools call LlmResponse.Builder responseBuilder = LlmResponse.builder(); - List parts = new ArrayList<>(); - Part part = ollamaContentBlockToPart(responseQuantum); - parts.add(part); - - // Call tool - if (!part.functionCall().isEmpty() - && part.functionResponse().isEmpty() - && !LAST_RESP_TOOl_EXECUTED) { - - responseBuilder.content( - Content.builder() - .role("assistant") - .parts( - ImmutableList.of(Part.builder().functionCall(part.functionCall().get()).build())) - .build()); - - // responseBuilder.partial(false).turnComplete(false); - - } else { - responseBuilder.content( - Content.builder().role("assistant").parts(ImmutableList.copyOf(parts)).build()); + List parts = ollamaContentBlockToParts(responseQuantum); + + responseBuilder.content( + Content.builder().role("assistant").parts(ImmutableList.copyOf(parts)).build()); + + if (usageMetadata != null) { + responseBuilder.usageMetadata(usageMetadata); } return Flowable.just(responseBuilder.build()); } + /** Best-effort extraction of a human readable error message from a Bedrock error body. */ + private static String extractBedrockErrorMessage(String body) { + if (body == null || body.isEmpty()) { + return null; + } + try { + JSONObject json = new JSONObject(body); + Object message = getKeyIgnoreCase(json, "message", "Message", "errorMessage", "__type"); + if (message != null) { + return message.toString(); + } + } catch (org.json.JSONException ignored) { + // Not JSON; fall through to raw body. + } + return body; + } + /** Gets a value from JSONObject using case-insensitive key match. */ private static Object getKeyIgnoreCase(JSONObject obj, String... keys) { Iterator it = obj.keys(); @@ -480,136 +482,15 @@ public Flowable generateContentStream(LlmRequest llmRequest) { } } - // Messages - JSONArray messages = new JSONArray(); - - JSONObject llmMessageJson1 = new JSONObject(); - llmMessageJson1.put("role", "assistant"); - JSONArray sysContentArray = new JSONArray(); - sysContentArray.put(new JSONObject().put("text", systemText)); - llmMessageJson1.put("content", sysContentArray); - messages.put(llmMessageJson1); // Agent system prompt is always added - - final List finalContents = contents; - finalContents.stream() - .forEach( - item -> { - JSONObject messageQuantum = new JSONObject(); - messageQuantum.put( - "role", - item.role().get().equals("model") || item.role().get().equals("assistant") - ? "assistant" - : "user"); - - // Additinal override work to add function response - if (item.parts().get().get(0).functionResponse().isPresent()) { - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put( - "text", - item.parts().get().get(0).functionResponse().get().name().get() - + " responded with these values, " - + new JSONObject( - item.parts().get().get(0).functionResponse().get().response().get()) - .toString()); - - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - - messageQuantum.put("content", contentArray2); - } else if (item.parts().get().get(0).functionCall().isPresent()) { - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put( - "text", - item.parts().get().get(0).functionCall().get().name().get() - + " is to be called with these arguments, " - + new JSONObject( - item.parts().get().get(0).functionCall().get().args().get()) - .toString()); - - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - - messageQuantum.put("content", contentArray2); - } else { - - JSONObject txtMsg3 = new JSONObject(); - txtMsg3.put("text", item.text()); - JSONArray contentArray2 = new JSONArray(); - contentArray2.put(txtMsg3); - messageQuantum.put("content", contentArray2); - } - messages.put(messageQuantum); - }); - - // Tools - JSONArray functions = new JSONArray(); - llmRequest - .tools() - .entrySet() - .forEach( - tooldetail -> { - BaseTool baseTool = tooldetail.getValue(); - Optional declarationOptional = baseTool.declaration(); - if (!declarationOptional.isPresent()) { - System.err.println( - "Skipping tool '" + baseTool.name() + "' with missing declaration."); - return; - } - FunctionDeclaration functionDeclaration = declarationOptional.get(); - Map toolMap = new HashMap<>(); - toolMap.put("name", cleanForIdentifierPattern(functionDeclaration.name().get())); - toolMap.put( - "description", - cleanForIdentifierPattern(functionDeclaration.description().orElse(""))); - Optional parametersOptional = functionDeclaration.parameters(); - if (parametersOptional.isPresent()) { - Schema parametersSchema = parametersOptional.get(); - Map parametersMap = new HashMap<>(); - parametersMap.put("type", "object"); - Optional> propertiesOptional = parametersSchema.properties(); - if (propertiesOptional.isPresent()) { - Map propertiesMap = new HashMap<>(); - ObjectMapper objectMapper = new ObjectMapper(); - objectMapper.registerModule(new Jdk8Module()); - propertiesOptional - .get() - .forEach( - (key, schema) -> { - Map schemaMap = - objectMapper.convertValue( - schema, new TypeReference>() {}); - updateTypeString(schemaMap); - propertiesMap.put(key, schemaMap); - }); - parametersMap.put("properties", propertiesMap); - } - parametersSchema - .required() - .ifPresent(requiredList -> parametersMap.put("required", requiredList)); - toolMap.put("parameters", parametersMap); - } - // Convert the complete tool map into an org.json.JSONObject - JSONObject jsonToolW = new JSONObject(); - jsonToolW.put("type", "function"); - - JSONObject jsonTool = new JSONObject(toolMap); - jsonToolW.put("function", jsonTool); + JSONArray system = buildSystemArray(systemText); + JSONArray messages = buildMessagesFromContents(contents); - // Add the generated tool JSON object to your functions list/array - functions.put(jsonToolW); - }); + JSONArray functions = buildToolsFromRequest(llmRequest); String modelId = this.model(); - boolean LAST_RESP_TOOl_EXECUTED = - Iterables.getLast(Iterables.getLast(finalContents).parts().get()) - .functionResponse() - .isPresent(); - return createRobustStreamingResponse( - modelId, - messages, - LAST_RESP_TOOl_EXECUTED ? null : (functions.length() > 0 ? functions : null)); + modelId, system, messages, functions.length() > 0 ? functions : null); } /** @@ -617,10 +498,11 @@ public Flowable generateContentStream(LlmRequest llmRequest) { * the "onNext already called in this generate turn" error. */ private Flowable createRobustStreamingResponse( - String modelId, JSONArray messages, JSONArray functions) { + String modelId, JSONArray system, JSONArray messages, JSONArray functions) { final StringBuilder accumulatedText = new StringBuilder(); final StringBuilder functionCallName = new StringBuilder(); final StringBuilder functionCallArgs = new StringBuilder(); + final StringBuilder functionCallToolUseId = new StringBuilder(); final AtomicBoolean inFunctionCall = new AtomicBoolean(false); final AtomicBoolean streamCompleted = new AtomicBoolean(false); @@ -632,7 +514,7 @@ private Flowable createRobustStreamingResponse( final AtomicInteger completionAudioTokens = new AtomicInteger(0); return Flowable.generate( - () -> callLLMChatStream(modelId, messages, functions), + () -> callLLMChatStream(modelId, system, messages, functions), (reader, emitter) -> { try { if (reader == null || streamCompleted.get()) { @@ -769,6 +651,22 @@ private Flowable createRobustStreamingResponse( hasContent = true; } } + if (c.has("toolUse")) { + inFunctionCall.set(true); + JSONObject toolUse = c.getJSONObject("toolUse"); + if (toolUse.has("name")) { + functionCallName.setLength(0); + functionCallName.append(toolUse.getString("name")); + } + if (toolUse.has("toolUseId")) { + functionCallToolUseId.setLength(0); + functionCallToolUseId.append(toolUse.getString("toolUseId")); + } + if (toolUse.has("input")) { + functionCallArgs.setLength(0); + functionCallArgs.append(toolUse.getJSONObject("input").toString()); + } + } } } @@ -803,7 +701,8 @@ private Flowable createRobustStreamingResponse( "end_turn".equals(stopReason) || "stop_sequence".equals(stopReason) || "max_tokens".equals(stopReason) - || "content_filtered".equals(stopReason); + || "content_filtered".equals(stopReason) + || "tool_use".equals(stopReason); } else if (responseJson.optBoolean("done", false)) { isDone = true; } @@ -824,16 +723,23 @@ private Flowable createRobustStreamingResponse( if (inFunctionCall.get() && functionCallName.length() > 0) { try { Map args = new JSONObject(functionCallArgs.toString()).toMap(); - FunctionCall fc = - FunctionCall.builder().name(functionCallName.toString()).args(args).build(); - Part part = Part.builder().functionCall(fc).build(); + FunctionCall.Builder fcBuilder = + FunctionCall.builder().name(functionCallName.toString()).args(args); + if (functionCallToolUseId.length() > 0) { + fcBuilder.id(functionCallToolUseId.toString()); + } + List responseParts = new ArrayList<>(); + if (accumulatedText.length() > 0) { + responseParts.add(Part.fromText(accumulatedText.toString())); + } + responseParts.add(Part.builder().functionCall(fcBuilder.build()).build()); LlmResponse.Builder functionResponseBuilder = LlmResponse.builder() .content( Content.builder() .role("model") - .parts(ImmutableList.of(part)) + .parts(ImmutableList.copyOf(responseParts)) .build()); if (usageMetadata != null) { @@ -917,7 +823,8 @@ private LlmResponse createTextResponse(String text, boolean partial) { .build(); } - public BufferedReader callLLMChatStream(String model, JSONArray messages, JSONArray tools) { + public BufferedReader callLLMChatStream( + String model, JSONArray system, JSONArray messages, JSONArray tools) { try { String bearerToken = getBearerToken(); if (bearerToken == null || bearerToken.isBlank()) { @@ -926,15 +833,9 @@ public BufferedReader callLLMChatStream(String model, JSONArray messages, JSONAr + "BEDROCK_BEARER_TOKEN, BEDROCK_API_KEY, BEDROCK_TOKEN (e.g. in .bashrc)"); } String baseUrl = getBedrockBaseUrl(D_URL); - String apiUrl = buildConverseUrl(baseUrl, model); + String apiUrl = buildConverseUrl(baseUrl, model).replace("/converse", "/converse-stream"); System.out.println("Using Bedrock URL: " + apiUrl); - JSONObject payload = new JSONObject(); - // Model already encoded in path; omit 'model' field to avoid Unexpected field type errors - payload.put("stream", true); - payload.put("messages", messages); - if (tools != null) { - payload.put("tools", tools); - } + JSONObject payload = buildConversePayload(system, messages, tools); String jsonString = payload.toString(); @@ -972,11 +873,15 @@ public BufferedReader callLLMChatStream(String model, JSONArray messages, JSONAr logger.error( "Bedrock streaming request failed: status={} body={}", responseCode, errorResponse); connection.disconnect(); - return null; + throw new LlmHttpException( + responseCode, + extractBedrockErrorMessage(errorResponse.toString()), + errorResponse.toString()); } } catch (IOException ex) { logger.error("Error in callLLMChatStream", ex); - return null; + throw new LlmHttpException( + -1, "Bedrock streaming request failed: " + ex.getMessage(), null, ex); } } @@ -990,7 +895,7 @@ public BaseLlmConnection connect(LlmRequest llmRequest) { * schemas, which is not directly related to sending chat messages to Ollama. You might consider * removing it if it's no longer needed. */ - private void updateTypeString(Map valueDict) { + private static void updateTypeString(Map valueDict) { if (valueDict == null) { return; } @@ -1015,42 +920,54 @@ private void updateTypeString(Map valueDict) { } public static Part ollamaContentBlockToPart(JSONObject blockJson) { + List parts = ollamaContentBlockToParts(blockJson); + return parts.get(0); + } + + /** Parses all text and toolUse blocks from a Bedrock Converse message object. */ + public static List ollamaContentBlockToParts(JSONObject blockJson) { + List parts = new ArrayList<>(); if (blockJson.has("content")) { JSONArray contentArray = blockJson.getJSONArray("content"); for (int i = 0; i < contentArray.length(); i++) { JSONObject tempObj = contentArray.getJSONObject(i); if (tempObj.has("text")) { - return Part.builder().text(tempObj.getString("text")).build(); - } - if (tempObj.has("toolUse")) { + String text = tempObj.getString("text"); + if (text != null && !text.isEmpty()) { + parts.add(Part.builder().text(text).build()); + } + } else if (tempObj.has("toolUse")) { JSONObject toolUse = tempObj.getJSONObject("toolUse"); - if (toolUse != null) { - if (toolUse.has("name")) { - JSONObject input = toolUse.optJSONObject("input"); - Map args = input.toMap(); - FunctionCall functionCall = - FunctionCall.builder().name(toolUse.getString("name")).args(args).build(); - return Part.builder().functionCall(functionCall).build(); + if (toolUse != null && toolUse.has("name")) { + JSONObject input = toolUse.optJSONObject("input"); + Map args = input != null ? input.toMap() : new HashMap<>(); + FunctionCall.Builder fcBuilder = + FunctionCall.builder().name(toolUse.getString("name")).args(args); + if (toolUse.has("toolUseId")) { + fcBuilder.id(toolUse.getString("toolUseId")); } + parts.add(Part.builder().functionCall(fcBuilder.build()).build()); } } } } - throw new UnsupportedOperationException( - "Unsupported content block format or missing required fields: " + blockJson.toString()); + if (parts.isEmpty()) { + throw new UnsupportedOperationException( + "Unsupported content block format or missing required fields: " + blockJson.toString()); + } + return parts; } /** - * Makes a POST request to a specified URL with a dynamic JSON body. Fetches username and password - * from environment variables. + * POST to the Bedrock Converse API. * - * @param model - * @param messages The list of messages for the "request.messages" field. - * @param tools - * @return The response body as a String. - * @throws RuntimeException If environment variables are not set or JSON creation fails. + * @param model deployment / model id + * @param system system prompt blocks + * @param messages conversation messages + * @param tools tool specs (toolConfig.tools) */ - public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) { + public JSONObject callLLMChat( + String model, JSONArray system, JSONArray messages, JSONArray tools) { try { String bearerToken = getBearerToken(); if (bearerToken == null || bearerToken.isBlank()) { @@ -1059,15 +976,8 @@ public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) + "BEDROCK_BEARER_TOKEN, BEDROCK_API_KEY, BEDROCK_TOKEN (e.g. in .bashrc)"); } String baseUrl = getBedrockBaseUrl(D_URL); - JSONObject responseJ = new JSONObject(); String apiUrl = buildConverseUrl(baseUrl, model); - JSONObject payload = new JSONObject(); - // Model already in path; omit from payload to avoid "Unexpected field type" errors - payload.put("stream", false); - payload.put("messages", messages); - if (tools != null) { - payload.put("tools", tools); - } + JSONObject payload = buildConversePayload(system, messages, tools); String jsonString = payload.toString(); URL url = new URL(apiUrl); @@ -1083,8 +993,6 @@ public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) writer.write(jsonString); System.out.println("Bedrock Base LM => " + jsonString); writer.flush(); - } catch (IOException ex) { - java.util.logging.Logger.getLogger(RedbusADG.class.getName()).log(Level.SEVERE, null, ex); } int responseCode = connection.getResponseCode(); @@ -1092,6 +1000,7 @@ public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) InputStream inputStream = (responseCode < 400) ? connection.getInputStream() : connection.getErrorStream(); + String raw = ""; try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, "UTF-8"))) { StringBuilder response = new StringBuilder(); @@ -1099,21 +1008,41 @@ public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) while ((line = reader.readLine()) != null) { response.append(line); } - responseJ = new JSONObject(response.toString()); - } catch (IOException ex) { - java.util.logging.Logger.getLogger(RedbusADG.class.getName()).log(Level.SEVERE, null, ex); + raw = response.toString(); } connection.disconnect(); - return responseJ; - } catch (MalformedURLException ex) { - java.util.logging.Logger.getLogger(RedbusADG.class.getName()).log(Level.SEVERE, null, ex); - } catch (ProtocolException ex) { - java.util.logging.Logger.getLogger(RedbusADG.class.getName()).log(Level.SEVERE, null, ex); + if (responseCode >= 400) { + throw new LlmHttpException(responseCode, extractBedrockErrorMessage(raw), raw); + } + try { + return new JSONObject(raw); + } catch (org.json.JSONException ex) { + throw new LlmHttpException(responseCode, "Bedrock returned a non-JSON response", raw); + } + + } catch (LlmHttpException e) { + throw e; + } catch (MalformedURLException | ProtocolException ex) { + throw new LlmHttpException(-1, ex.getMessage(), null, ex); } catch (IOException ex) { - java.util.logging.Logger.getLogger(RedbusADG.class.getName()).log(Level.SEVERE, null, ex); + throw new LlmHttpException(-1, "Bedrock request failed: " + ex.getMessage(), null, ex); } - return new JSONObject(); + } + + /** + * Makes a POST request to a specified URL with a dynamic JSON body. Fetches username and password + * from environment variables. + * + * @param model + * @param messages The list of messages for the "request.messages" field. + * @param tools + * @return The response body as a String. + * @throws RuntimeException If environment variables are not set or JSON creation fails. + */ + @Deprecated + public JSONObject callLLMChat(String model, JSONArray messages, JSONArray tools) { + return callLLMChat(model, buildSystemArray(""), messages, tools); } /** @@ -1410,34 +1339,19 @@ private GenerateContentResponseUsageMetadata getUsageMetadata(JSONObject agentRe } public static void main(String[] args) { - // --- Create the 'messages' part of the JSON using org.json --- - String messagesJsonString = - """ - [ - { - "role": "assistant", - "content": [ { "text": "You are a helpful assistant." } ] - }, - { - "role": "user", - "content": [ { "text": "Write a story about a curious cat named Tommy who explores a mysterious garden. Make the story at least 10 lines long, with each line describing a new discovery or adventure Tommy has in the garden. End with Tommy finding a new friend." } ] - } - ] - """; - - JSONArray messagesArray; - try { - messagesArray = new JSONArray(messagesJsonString); - } catch (Exception e) { - System.err.println("Failed to parse JSON string into JSONArray: " + e.getMessage()); - return; - } - - String modelId = "openai.gpt-oss-120b-1:0"; // Example model ID for Bedrock - String bedrockBaseUrl = - "https://bedrock-runtime.us-west-2.amazonaws.com/model"; // Base URL only + String modelId = "openai.gpt-oss-120b-1:0"; + String bedrockBaseUrl = "https://bedrock-runtime.us-west-2.amazonaws.com/model"; BedrockBaseLM ollamaLlm = new BedrockBaseLM(modelId, bedrockBaseUrl); + JSONArray systemArray = buildSystemArray("You are a helpful assistant."); + JSONArray userOnlyMessages = new JSONArray(); + userOnlyMessages.put( + new JSONObject() + .put("role", "user") + .put( + "content", + new JSONArray().put(new JSONObject().put("text", "Say hello in one word.")))); + // --- Test Streaming Call --- System.out.println("--- Testing Streaming API Call ---"); try { @@ -1445,7 +1359,8 @@ public static void main(String[] args) { System.out.println("Using model ID: " + modelId); System.out.println("Fetching Ollama endpoint from environment variable: " + BEDROCK_ENV_VAR); - BufferedReader responseReader = ollamaLlm.callLLMChatStream(modelId, messagesArray, null); + BufferedReader responseReader = + ollamaLlm.callLLMChatStream(modelId, systemArray, userOnlyMessages, null); if (responseReader != null) { System.out.println("\nAPI Call Successful! Streaming response:"); @@ -1511,7 +1426,7 @@ public static void main(String[] args) { System.out.println("Attempting to call Ollama API (Non-Streaming)..."); System.out.println("Using model ID: " + modelId); - JSONObject responseJson = ollamaLlm.callLLMChat(modelId, messagesArray, null); + JSONObject responseJson = ollamaLlm.callLLMChat(modelId, systemArray, userOnlyMessages, null); if (responseJson != null && !responseJson.isEmpty()) { System.out.println("\nAPI Call Successful! Non-Streaming response:"); diff --git a/core/src/main/java/com/google/adk/models/Claude.java b/core/src/main/java/com/google/adk/models/Claude.java index fe8a4f2e5..07f5f8fde 100644 --- a/core/src/main/java/com/google/adk/models/Claude.java +++ b/core/src/main/java/com/google/adk/models/Claude.java @@ -36,6 +36,7 @@ import com.google.common.collect.ImmutableMap; import com.google.genai.types.*; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -151,11 +152,17 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre paramsBuilder.toolChoice(toolChoice); } - var message = this.anthropicClient.messages().create(paramsBuilder.build()); - - logger.debug("Claude response: {}", message); - - return Flowable.just(convertAnthropicResponseToLlmResponse(message)); + MessageCreateParams params = paramsBuilder.build(); + + // Defer the blocking Anthropic SDK call until subscription so it is cancellable and bounded by + // any downstream timeout, enabling latency-based failover in FailoverLlm. + return Flowable.defer( + () -> { + var message = this.anthropicClient.messages().create(params); + logger.debug("Claude response: {}", message); + return Flowable.just(convertAnthropicResponseToLlmResponse(message)); + }) + .subscribeOn(Schedulers.io()); } private Role toClaudeRole(String role) { diff --git a/core/src/main/java/com/google/adk/models/LlmRegistry.java b/core/src/main/java/com/google/adk/models/LlmRegistry.java index 8c0b5d24b..69ccc2565 100644 --- a/core/src/main/java/com/google/adk/models/LlmRegistry.java +++ b/core/src/main/java/com/google/adk/models/LlmRegistry.java @@ -16,10 +16,15 @@ package com.google.adk.models; +import com.google.adk.models.failover.FailoverConfig; +import com.google.adk.models.failover.FailoverLlm; import com.google.common.annotations.VisibleForTesting; +import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; /** Central registry for managing Large Language Model (LLM) instances. */ public final class LlmRegistry { @@ -41,6 +46,16 @@ public interface LlmFactory { */ private static final LinkedHashMap llmFactories = new LinkedHashMap<>(); + /** + * Regex patterns to ordered fallback model-name chains, in registration order. When a resolved + * model name matches one of these patterns, {@link #getLlm} wraps the created model in a {@link + * FailoverLlm} that fails over to the listed models. First match wins. + */ + private static final LinkedHashMap> fallbackChains = new LinkedHashMap<>(); + + /** Config applied to {@link FailoverLlm} instances built by the registry. */ + private static volatile FailoverConfig defaultFailoverConfig = FailoverConfig.defaults(); + /** Registers default LLM factories, e.g. for Gemini models. */ static { registerLlm("gemini-.*", modelName -> Gemini.builder().modelName(modelName).build()); @@ -95,7 +110,85 @@ static boolean matchesAnyPattern(String modelName) { * @throws IllegalArgumentException If no factory matches the model name. */ public static BaseLlm getLlm(String modelName) { - return instances.computeIfAbsent(modelName, LlmRegistry::createLlm); + return instances.computeIfAbsent(modelName, LlmRegistry::createLlmWithFallback); + } + + /** + * Registers an ordered fallback chain for model names matching the given regex pattern. Models + * resolved through {@link #getLlm} whose name matches the pattern are wrapped in a {@link + * FailoverLlm} that tries the matched model first, then each fallback model in order on a + * classified, retryable failure or timeout. + * + *

Fallback model names are resolved lazily (only on failure) through the registered factories, + * so a chain may mix providers, e.g. {@code ["gemini-2.5-flash-lite", "Bedrock|...claude..."]}. + * + * @param modelNamePattern Regex pattern for matching primary model names. + * @param fallbackModelNames Ordered fallback model names (excluding the primary). + */ + public static void registerFallback(String modelNamePattern, List fallbackModelNames) { + synchronized (factoryLock) { + fallbackChains.put(modelNamePattern, new ArrayList<>(fallbackModelNames)); + } + // Drop cached instances so subsequent lookups pick up the new (un)wrapped form. + instances.keySet().removeIf(name -> name.matches(modelNamePattern)); + } + + /** Sets the {@link FailoverConfig} used for registry-built {@link FailoverLlm} instances. */ + public static void setDefaultFailoverConfig(FailoverConfig config) { + defaultFailoverConfig = config == null ? FailoverConfig.defaults() : config; + // Existing wrapped instances captured the old config; drop them so they rebuild. + if (!fallbackChains.isEmpty()) { + instances.clear(); + } + } + + /** + * Evicts cached LLM instances whose model name matches the given regex pattern. Use this to + * switch a model at runtime: change the factory/config, then evict so the next {@link #getLlm} + * rebuilds the instance. + * + * @param modelNamePattern Regex pattern for matching cached model names. + */ + public static void evict(String modelNamePattern) { + instances.keySet().removeIf(name -> name.matches(modelNamePattern)); + } + + /** Clears all cached LLM instances. The next {@link #getLlm} call rebuilds them. */ + public static void clearInstances() { + instances.clear(); + } + + private static BaseLlm createLlmWithFallback(String modelName) { + BaseLlm primary = createLlm(modelName); + List fallbacks = findFallbackChain(modelName); + if (fallbacks == null || fallbacks.isEmpty()) { + return primary; + } + List> suppliers = new ArrayList<>(); + suppliers.add(() -> primary); + for (String fallbackName : fallbacks) { + if (fallbackName == null || fallbackName.isEmpty() || fallbackName.equals(modelName)) { + continue; + } + // Resolve fallbacks via the raw factory path (not getLlm) to avoid re-wrapping/recursion; + // FailoverLlm memoizes each supplier so the instance is built at most once, only on failure. + suppliers.add(() -> createLlm(fallbackName)); + } + if (suppliers.size() == 1) { + return primary; + } + return new FailoverLlm(suppliers, defaultFailoverConfig); + } + + private static List findFallbackChain(String modelName) { + synchronized (factoryLock) { + for (Map.Entry> entry : fallbackChains.entrySet()) { + if (modelName.matches(entry.getKey())) { + return entry.getValue(); + } + } + } + return null; } /** diff --git a/core/src/main/java/com/google/adk/models/azure/AzureConfig.java b/core/src/main/java/com/google/adk/models/azure/AzureConfig.java index a162c4f86..a38bfbc1f 100644 --- a/core/src/main/java/com/google/adk/models/azure/AzureConfig.java +++ b/core/src/main/java/com/google/adk/models/azure/AzureConfig.java @@ -1,6 +1,8 @@ package com.google.adk.models.azure; import java.util.Collection; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -26,10 +28,13 @@ public final class AzureConfig { private static final Logger logger = LoggerFactory.getLogger(AzureConfig.class); + public static final String MODEL_ENDPOINT_ENV = "AZURE_MODEL_ENDPOINT"; public static final String RESPONSE_ENDPOINT_ENV = "AZURE_RESPONSE_ENDPOINT"; public static final String REALTIME_ENDPOINT_ENV = "AZURE_REALTIME_ENDPOINT"; public static final String TRANSLATE_ENDPOINT_ENV = "AZURE_TRANSLATE_ENDPOINT"; + private static final Map runtimeOverrides = new ConcurrentHashMap<>(); + public static final String API_KEY_ENV = "AZURE_OPENAI_API_KEY"; public static final String VOICE_ENV = "AZURE_REALTIME_VOICE"; public static final String TRANSLATE_TARGET_LANGUAGE_ENV = "AZURE_TRANSLATE_TARGET_LANGUAGE"; @@ -82,6 +87,17 @@ private AzureConfig( this.translateTargetLanguage = translateTargetLanguage; } + /** + * Sets a runtime configuration value that takes precedence over environment variables. Used by + * host applications (e.g. RAE) to inject values from {@code application.properties} or AWS + * Secrets Manager at startup. + */ + public static void setRuntimeOverride(String envKey, String value) { + if (value != null && !value.isBlank()) { + runtimeOverrides.put(envKey, value.replaceAll("/+$", "")); + } + } + public static AzureConfig fromEnvironment(String modelName) { String responseEndpoint = resolveContractEndpoint(RESPONSE_ENDPOINT_ENV, "Responses API"); String realtimeEndpoint = resolveContractEndpoint(REALTIME_ENDPOINT_ENV, "Realtime voice API"); @@ -217,6 +233,9 @@ static String normalizeTranslateWebSocketUrl(String raw, String modelName) { private static String resolveContractEndpoint(String specificEnv, String label) { String val = resolveOptionalEnv(specificEnv); + if ((val == null || val.isBlank()) && RESPONSE_ENDPOINT_ENV.equals(specificEnv)) { + val = resolveOptionalEnv(MODEL_ENDPOINT_ENV); + } if (val == null || val.isBlank()) { throw new IllegalStateException( "Azure " + label + " endpoint not configured. Set " + specificEnv); @@ -256,7 +275,7 @@ private static String toWebSocketUrl(String url) { } private static String resolveRequired(String envVar) { - String val = System.getenv(envVar); + String val = resolveEnvOrOverride(envVar); if (val == null || val.isBlank()) { throw new IllegalStateException(envVar + " environment variable is not set."); } @@ -264,15 +283,23 @@ private static String resolveRequired(String envVar) { } private static String resolveOptional(String envVar, String defaultValue) { - String val = System.getenv(envVar); + String val = resolveEnvOrOverride(envVar); return (val != null && !val.isBlank()) ? val : defaultValue; } private static String resolveOptionalEnv(String envVar) { - String val = System.getenv(envVar); + String val = resolveEnvOrOverride(envVar); return (val != null && !val.isBlank()) ? val.replaceAll("/+$", "") : null; } + private static String resolveEnvOrOverride(String envVar) { + String override = runtimeOverrides.get(envVar); + if (override != null && !override.isBlank()) { + return override; + } + return System.getenv(envVar); + } + private static String maskEndpoint(String url) { if (url == null) { return "unset"; diff --git a/core/src/main/java/com/google/adk/models/azure/AzureRestTransport.java b/core/src/main/java/com/google/adk/models/azure/AzureRestTransport.java index ff244b17d..d5ca7f021 100644 --- a/core/src/main/java/com/google/adk/models/azure/AzureRestTransport.java +++ b/core/src/main/java/com/google/adk/models/azure/AzureRestTransport.java @@ -7,6 +7,7 @@ import com.google.adk.models.GenericLlmConnection; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.models.failover.LlmHttpException; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.genai.types.Content; @@ -15,6 +16,7 @@ import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; @@ -95,49 +97,61 @@ public BaseLlmConnection connect(LlmRequest req) { // ==================== Non-streaming ==================== private Flowable generateContentSync(LlmRequest llmRequest, AzureConfig config) { - List contents = ensureLastContentIsUser(llmRequest.contents()); - String instructions = AzureRequestConverter.extractInstructions(llmRequest); - JSONArray inputItems = buildInputItems(contents); - JSONArray tools = AzureRequestConverter.buildTools(llmRequest); - - boolean lastRespToolExecuted = - Iterables.getLast(Iterables.getLast(contents).parts().get()).functionResponse().isPresent(); - - Optional temperature = llmRequest.config().flatMap(GenerateContentConfig::temperature); - Optional maxTokens = - llmRequest.config().flatMap(GenerateContentConfig::maxOutputTokens); - - JSONObject payload = new JSONObject(); - payload.put("model", config.modelName()); - payload.put("input", inputItems); - if (!instructions.isEmpty()) { - payload.put("instructions", instructions); - } - temperature.ifPresent(t -> payload.put("temperature", t)); - payload.put("stream", false); - payload.put("store", false); - payload.put("reasoning", new JSONObject().put("summary", "auto")); - if (maxTokens.isPresent() && maxTokens.get() > 0) { - payload.put("max_output_tokens", maxTokens.get()); - } - if (!lastRespToolExecuted && tools.length() > 0) { - payload.put("tools", tools); - } + // Defer the blocking HTTP call until subscription so it is cancellable and can be bounded by a + // downstream timeout (enabling latency-based failover in FailoverLlm). + return Flowable.defer( + () -> { + List contents = ensureLastContentIsUser(llmRequest.contents()); + String instructions = AzureRequestConverter.extractInstructions(llmRequest); + JSONArray inputItems = buildInputItems(contents); + JSONArray tools = AzureRequestConverter.buildTools(llmRequest); + + boolean lastRespToolExecuted = + Iterables.getLast(Iterables.getLast(contents).parts().get()) + .functionResponse() + .isPresent(); + + Optional temperature = + llmRequest.config().flatMap(GenerateContentConfig::temperature); + Optional maxTokens = + llmRequest.config().flatMap(GenerateContentConfig::maxOutputTokens); + + JSONObject payload = new JSONObject(); + payload.put("model", config.modelName()); + payload.put("input", inputItems); + if (!instructions.isEmpty()) { + payload.put("instructions", instructions); + } + temperature.ifPresent(t -> payload.put("temperature", t)); + payload.put("stream", false); + payload.put("store", false); + payload.put("reasoning", new JSONObject().put("summary", "auto")); + if (maxTokens.isPresent() && maxTokens.get() > 0) { + payload.put("max_output_tokens", maxTokens.get()); + } + if (!lastRespToolExecuted && tools.length() > 0) { + payload.put("tools", tools); + } - logger.debug("Azure Responses API request payload size: {} bytes", payload.toString().length()); + logger.debug( + "Azure Responses API request payload size: {} bytes", + payload.toString().length()); - JSONObject response = callApi(payload, config); + JSONObject response = callApi(payload, config); - if (response.has("error") && !response.isNull("error")) { - JSONObject error = response.getJSONObject("error"); - String message = error.optString("message", response.toString()); - logger.error("Azure Responses API error: {}", response); - return Flowable.error(new IllegalStateException("Azure Responses API error: " + message)); - } + if (response.has("error") && !response.isNull("error")) { + JSONObject error = response.getJSONObject("error"); + String message = error.optString("message", response.toString()); + logger.error("Azure Responses API error: {}", response); + return Flowable.error( + new IllegalStateException("Azure Responses API error: " + message)); + } - GenerateContentResponseUsageMetadata usageMetadata = extractUsageMetadata(response); - LlmResponse llmResponse = parseOutputToLlmResponse(response, usageMetadata); - return Flowable.just(llmResponse); + GenerateContentResponseUsageMetadata usageMetadata = extractUsageMetadata(response); + LlmResponse llmResponse = parseOutputToLlmResponse(response, usageMetadata); + return Flowable.just(llmResponse); + }) + .subscribeOn(Schedulers.io()); } // ==================== Streaming ==================== @@ -185,304 +199,313 @@ private Flowable generateContentStream(LlmRequest llmRequest, Azure logger.debug("Starting streaming request for model: {}", config.modelName()); logger.debug("Streaming payload size: {} bytes", payload.toString().length()); - return Flowable.create( - emitter -> { - BufferedReader reader = null; - try { - logger.debug("Opening SSE connection..."); - reader = callApiStream(payload, config); - if (reader == null) { - logger.warn("Azure SSE reader is null — stream failed to open."); - emitter.onComplete(); - return; - } - logger.debug("SSE connection opened successfully."); - long streamStartMs = System.currentTimeMillis(); - int chunkCount = 0; - - String lastEventName = null; - String line; - while ((line = reader.readLine()) != null) { - if (emitter.isCancelled()) { - logger.debug("Emitter cancelled, breaking out of read loop."); - break; - } - - logger.debug( - "SSE raw: {}", line.length() > 200 ? line.substring(0, 200) + "..." : line); - - if (line.isEmpty()) continue; - if (line.startsWith("event:")) { - lastEventName = line.substring(6).trim(); - continue; - } - if (!line.startsWith("data:")) continue; - - String jsonStr = line.substring(5).trim(); - if (jsonStr.equals("[DONE]")) { - long elapsed = System.currentTimeMillis() - streamStartMs; - logger.debug( - "[DONE] marker received after {}ms, total chunks: {}", elapsed, chunkCount); - break; - } - - chunkCount++; - JSONObject event; + return Flowable.create( + emitter -> { + BufferedReader reader = null; try { - event = new JSONObject(jsonStr); - } catch (JSONException e) { - logger.warn("Failed to parse SSE chunk #{}: {}", chunkCount, jsonStr); - continue; - } - - String eventType = event.optString("type", ""); - if (eventType.isEmpty() && lastEventName != null) { - eventType = lastEventName; - } - lastEventName = null; - - logger.debug( - "SSE chunk #{} eventType='{}' keys={}", chunkCount, eventType, event.keySet()); - - switch (eventType) { - case "response.output_item.added": - { - JSONObject item = event.optJSONObject("item"); - if (item == null) break; - String itemType = item.optString("type", ""); - if ("function_call".equals(itemType)) { - inFunctionCall.set(true); - String name = item.optString("name", ""); - String callId = item.optString("call_id", ""); - logger.debug("Function call starting: name='{}' callId='{}'", name, callId); - if (!name.isEmpty()) functionCallName.append(name); - if (!callId.isEmpty()) functionCallCallId.append(callId); - } else if ("reasoning".equals(itemType)) { - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText("\ud83e\udde0 Thinking...\n")) - .build()) - .partial(true) - .build()); - } + logger.debug("Opening SSE connection..."); + reader = callApiStream(payload, config); + if (reader == null) { + logger.warn("Azure SSE reader is null — stream failed to open."); + if (!emitter.isCancelled()) { + emitter.onError( + new LlmHttpException(-1, "Azure SSE stream failed to open", null)); + } + return; + } + logger.debug("SSE connection opened successfully."); + long streamStartMs = System.currentTimeMillis(); + int chunkCount = 0; + + String lastEventName = null; + String line; + while ((line = reader.readLine()) != null) { + if (emitter.isCancelled()) { + logger.debug("Emitter cancelled, breaking out of read loop."); break; } - case "response.reasoning_summary_text.delta": - { - String delta = event.optString("delta", ""); - if (!delta.isEmpty()) { - reasoningSummary.append(delta); - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText(delta)) - .build()) - .partial(true) - .build()); - } - break; + logger.debug( + "SSE raw: {}", line.length() > 200 ? line.substring(0, 200) + "..." : line); + + if (line.isEmpty()) continue; + if (line.startsWith("event:")) { + lastEventName = line.substring(6).trim(); + continue; } + if (!line.startsWith("data:")) continue; - case "response.reasoning_summary_text.done": - { - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText("\n\n")) - .build()) - .partial(true) - .build()); + String jsonStr = line.substring(5).trim(); + if (jsonStr.equals("[DONE]")) { + long elapsed = System.currentTimeMillis() - streamStartMs; + logger.debug( + "[DONE] marker received after {}ms, total chunks: {}", elapsed, chunkCount); break; } - case "response.output_text.delta": - { - String delta = extractTextDeltaFromStreamEvent(event); - if (!delta.isEmpty()) { - accumulatedText.append(delta); - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText(delta)) - .build()) - .partial(true) - .build()); - } - break; + chunkCount++; + JSONObject event; + try { + event = new JSONObject(jsonStr); + } catch (JSONException e) { + logger.warn("Failed to parse SSE chunk #{}: {}", chunkCount, jsonStr); + continue; } - case "response.output_text.done": - { - String fullText = event.optString("text", ""); - if (!fullText.isEmpty()) { - accumulatedText.setLength(0); - accumulatedText.append(fullText); - finalTextEmitted.set(true); - String finalContent = fullText; - if (reasoningSummary.length() > 0) { - finalContent = - "\ud83e\udde0 **Thinking:**\n> " - + reasoningSummary.toString().replace("\n", "\n> ") - + "\n\n" - + fullText; - } - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText(finalContent)) - .build()) - .partial(false) - .build()); - } - break; + String eventType = event.optString("type", ""); + if (eventType.isEmpty() && lastEventName != null) { + eventType = lastEventName; } + lastEventName = null; + + logger.debug( + "SSE chunk #{} eventType='{}' keys={}", + chunkCount, + eventType, + event.keySet()); + + switch (eventType) { + case "response.output_item.added": + { + JSONObject item = event.optJSONObject("item"); + if (item == null) break; + String itemType = item.optString("type", ""); + if ("function_call".equals(itemType)) { + inFunctionCall.set(true); + String name = item.optString("name", ""); + String callId = item.optString("call_id", ""); + logger.debug( + "Function call starting: name='{}' callId='{}'", name, callId); + if (!name.isEmpty()) functionCallName.append(name); + if (!callId.isEmpty()) functionCallCallId.append(callId); + } else if ("reasoning".equals(itemType)) { + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts(Part.fromText("\ud83e\udde0 Thinking...\n")) + .build()) + .partial(true) + .build()); + } + break; + } - case "response.output_item.done": - { - if (finalTextEmitted.get()) break; - JSONObject item = event.optJSONObject("item"); - if (item != null && "message".equals(item.optString("type"))) { - String fullText = extractTextFromOutputMessageItem(item); - if (!fullText.isEmpty()) { - accumulatedText.setLength(0); - accumulatedText.append(fullText); - finalTextEmitted.set(true); - String finalContent = fullText; - if (reasoningSummary.length() > 0) { - finalContent = - "\ud83e\udde0 **Thinking:**\n> " - + reasoningSummary.toString().replace("\n", "\n> ") - + "\n\n" - + fullText; + case "response.reasoning_summary_text.delta": + { + String delta = event.optString("delta", ""); + if (!delta.isEmpty()) { + reasoningSummary.append(delta); + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts(Part.fromText(delta)) + .build()) + .partial(true) + .build()); } + break; + } + + case "response.reasoning_summary_text.done": + { emitter.onNext( LlmResponse.builder() .content( Content.builder() .role("model") - .parts(Part.fromText(finalContent)) + .parts(Part.fromText("\n\n")) .build()) - .partial(false) + .partial(true) .build()); + break; } - } - break; - } - case "response.function_call_arguments.delta": - { - String delta = extractTextDeltaFromStreamEvent(event); - if (!delta.isEmpty()) { - functionCallArgs.append(delta); - } - break; - } + case "response.output_text.delta": + { + String delta = extractTextDeltaFromStreamEvent(event); + if (!delta.isEmpty()) { + accumulatedText.append(delta); + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts(Part.fromText(delta)) + .build()) + .partial(true) + .build()); + } + break; + } - case "response.function_call_arguments.done": - { - if (functionCallName.length() > 0) { - String argsStr = - functionCallArgs.length() > 0 ? functionCallArgs.toString() : "{}"; - Map args; - try { - args = new JSONObject(argsStr).toMap(); - } catch (JSONException e) { - logger.warn("Failed to parse function args: {}", argsStr); - args = Map.of(); + case "response.output_text.done": + { + String fullText = event.optString("text", ""); + if (!fullText.isEmpty()) { + accumulatedText.setLength(0); + accumulatedText.append(fullText); + finalTextEmitted.set(true); + String finalContent = fullText; + if (reasoningSummary.length() > 0) { + finalContent = + "\ud83e\udde0 **Thinking:**\n> " + + reasoningSummary.toString().replace("\n", "\n> ") + + "\n\n" + + fullText; + } + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts(Part.fromText(finalContent)) + .build()) + .partial(false) + .build()); + } + break; } - FunctionCall fc = - FunctionCall.builder() - .name(functionCallName.toString()) - .args(args) - .build(); - emitter.onNext( - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts( - ImmutableList.of(Part.builder().functionCall(fc).build())) - .build()) - .partial(false) - .build()); - } - break; - } - case "response.completed": - { - JSONObject resp = event.optJSONObject("response"); - if (resp != null) { - JSONObject usage = resp.optJSONObject("usage"); - if (usage != null) { - inputTokens.set(usage.optInt("input_tokens", 0)); - outputTokens.set(usage.optInt("output_tokens", 0)); - logger.debug( - "Stream token usage — input: {}, output: {}", - inputTokens.get(), - outputTokens.get()); + case "response.output_item.done": + { + if (finalTextEmitted.get()) break; + JSONObject item = event.optJSONObject("item"); + if (item != null && "message".equals(item.optString("type"))) { + String fullText = extractTextFromOutputMessageItem(item); + if (!fullText.isEmpty()) { + accumulatedText.setLength(0); + accumulatedText.append(fullText); + finalTextEmitted.set(true); + String finalContent = fullText; + if (reasoningSummary.length() > 0) { + finalContent = + "\ud83e\udde0 **Thinking:**\n> " + + reasoningSummary.toString().replace("\n", "\n> ") + + "\n\n" + + fullText; + } + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts(Part.fromText(finalContent)) + .build()) + .partial(false) + .build()); + } + } + break; } - } - break; - } - default: - break; - } - } + case "response.function_call_arguments.delta": + { + String delta = extractTextDeltaFromStreamEvent(event); + if (!delta.isEmpty()) { + functionCallArgs.append(delta); + } + break; + } - long totalElapsed = System.currentTimeMillis() - streamStartMs; - logger.debug( - "Stream read loop finished — elapsed: {}ms, chunks: {}, accumulatedText: {} chars," - + " finalTextEmitted: {}, inFunctionCall: {}", - totalElapsed, - chunkCount, - accumulatedText.length(), - finalTextEmitted.get(), - inFunctionCall.get()); - - if (!emitter.isCancelled()) { - if (!finalTextEmitted.get()) { - emitFinalStreamResponse( - emitter, - accumulatedText, - inFunctionCall, - functionCallName, - functionCallArgs, - inputTokens.get(), - outputTokens.get()); - } - emitter.onComplete(); - } - } catch (IOException e) { - logger.error("IOException in Azure stream", e); - if (!emitter.isCancelled()) emitter.onError(e); - } catch (Exception e) { - logger.error("Error in Azure streaming", e); - if (!emitter.isCancelled()) emitter.onError(e); - } finally { - if (reader != null) { - try { - reader.close(); + case "response.function_call_arguments.done": + { + if (functionCallName.length() > 0) { + String argsStr = + functionCallArgs.length() > 0 ? functionCallArgs.toString() : "{}"; + Map args; + try { + args = new JSONObject(argsStr).toMap(); + } catch (JSONException e) { + logger.warn("Failed to parse function args: {}", argsStr); + args = Map.of(); + } + FunctionCall fc = + FunctionCall.builder() + .name(functionCallName.toString()) + .args(args) + .build(); + emitter.onNext( + LlmResponse.builder() + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder().functionCall(fc).build())) + .build()) + .partial(false) + .build()); + } + break; + } + + case "response.completed": + { + JSONObject resp = event.optJSONObject("response"); + if (resp != null) { + JSONObject usage = resp.optJSONObject("usage"); + if (usage != null) { + inputTokens.set(usage.optInt("input_tokens", 0)); + outputTokens.set(usage.optInt("output_tokens", 0)); + logger.debug( + "Stream token usage — input: {}, output: {}", + inputTokens.get(), + outputTokens.get()); + } + } + break; + } + + default: + break; + } + } + + long totalElapsed = System.currentTimeMillis() - streamStartMs; + logger.debug( + "Stream read loop finished — elapsed: {}ms, chunks: {}, accumulatedText: {} chars," + + " finalTextEmitted: {}, inFunctionCall: {}", + totalElapsed, + chunkCount, + accumulatedText.length(), + finalTextEmitted.get(), + inFunctionCall.get()); + + if (!emitter.isCancelled()) { + if (!finalTextEmitted.get()) { + emitFinalStreamResponse( + emitter, + accumulatedText, + inFunctionCall, + functionCallName, + functionCallArgs, + inputTokens.get(), + outputTokens.get()); + } + emitter.onComplete(); + } } catch (IOException e) { - logger.error("Error closing stream reader", e); + logger.error("IOException in Azure stream", e); + if (!emitter.isCancelled()) emitter.onError(e); + } catch (Exception e) { + logger.error("Error in Azure streaming", e); + if (!emitter.isCancelled()) emitter.onError(e); + } finally { + if (reader != null) { + try { + reader.close(); + } catch (IOException e) { + logger.error("Error closing stream reader", e); + } + } } - } - } - }, - io.reactivex.rxjava3.core.BackpressureStrategy.BUFFER); + }, + io.reactivex.rxjava3.core.BackpressureStrategy.BUFFER) + .subscribeOn(Schedulers.io()); } // ==================== Helpers ==================== @@ -628,22 +651,41 @@ private JSONObject callApi(JSONObject payload, AzureConfig config) { return new JSONObject(response.body()); } else { logger.error("Azure API error: status={} body={}", statusCode, response.body()); - try { - return new JSONObject(response.body()); - } catch (JSONException e) { - return new JSONObject().put("error", response.body()); - } + // Surface the status code as a typed exception so FailoverLlm can classify and fail over + // instead of silently returning an error body. + throw new LlmHttpException( + statusCode, extractAzureErrorMessage(response.body()), response.body()); } } catch (InterruptedException ex) { Thread.currentThread().interrupt(); logger.error("HTTP request interrupted for Azure Responses API", ex); - return new JSONObject().put("error", ex.getMessage()); + throw new LlmHttpException(-1, "Azure Responses API request interrupted", null, ex); } catch (IOException ex) { logger.error("HTTP request failed for Azure Responses API", ex); - return new JSONObject().put("error", ex.getMessage()); + throw new LlmHttpException( + -1, "Azure Responses API request failed: " + ex.getMessage(), null, ex); } } + private static String extractAzureErrorMessage(String body) { + if (body == null || body.isEmpty()) { + return null; + } + try { + JSONObject json = new JSONObject(body); + if (json.has("error") && !json.isNull("error")) { + Object error = json.get("error"); + if (error instanceof JSONObject) { + return ((JSONObject) error).optString("message", body); + } + return error.toString(); + } + } catch (JSONException ignored) { + // Not JSON; fall through to raw body. + } + return body; + } + private BufferedReader callApiStream(JSONObject payload, AzureConfig config) { try { String jsonString = payload.toString(); @@ -668,24 +710,27 @@ private BufferedReader callApiStream(JSONObject payload, AzureConfig config) { if (statusCode >= 200 && statusCode < 300) { return new BufferedReader(new InputStreamReader(response.body(), StandardCharsets.UTF_8)); } else { + StringBuilder errorBody = new StringBuilder(); try (BufferedReader errorReader = new BufferedReader(new InputStreamReader(response.body(), StandardCharsets.UTF_8))) { - StringBuilder errorBody = new StringBuilder(); String errorLine; while ((errorLine = errorReader.readLine()) != null) { errorBody.append(errorLine); } logger.error("Azure streaming failed: status={} body={}", statusCode, errorBody); } - return null; + // Surface the status so FailoverLlm can classify and fail over rather than complete empty. + throw new LlmHttpException( + statusCode, extractAzureErrorMessage(errorBody.toString()), errorBody.toString()); } } catch (InterruptedException ex) { Thread.currentThread().interrupt(); logger.error("HTTP request interrupted for Azure streaming", ex); - return null; + throw new LlmHttpException(-1, "Azure streaming request interrupted", null, ex); } catch (IOException ex) { logger.error("HTTP request failed for Azure streaming", ex); - return null; + throw new LlmHttpException( + -1, "Azure streaming request failed: " + ex.getMessage(), null, ex); } } diff --git a/core/src/main/java/com/google/adk/models/failover/FailoverConfig.java b/core/src/main/java/com/google/adk/models/failover/FailoverConfig.java new file mode 100644 index 000000000..09f813e57 --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/FailoverConfig.java @@ -0,0 +1,117 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +import java.time.Duration; +import java.util.Collections; +import java.util.EnumSet; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** Tuning knobs for {@link FailoverLlm}. */ +public final class FailoverConfig { + + private final Duration attemptTimeout; + private final boolean treatEmptyResponseAsFailure; + @Nullable private final Set retryableCategories; + + private FailoverConfig(Builder builder) { + this.attemptTimeout = builder.attemptTimeout; + this.treatEmptyResponseAsFailure = builder.treatEmptyResponseAsFailure; + this.retryableCategories = + builder.retryableCategories == null + ? null + : Collections.unmodifiableSet(EnumSet.copyOf(builder.retryableCategories)); + } + + /** Default config: no timeout, empty responses treated as failures, default retryable set. */ + public static FailoverConfig defaults() { + return builder().build(); + } + + public static Builder builder() { + return new Builder(); + } + + /** + * Per-attempt latency budget. {@link Duration#ZERO} (the default) disables timeout-based + * failover. When set, this bounds both time-to-first-response and the gap between streamed + * chunks; exceeding it raises a {@link java.util.concurrent.TimeoutException} which is classified + * as {@link LlmFailureCategory#TIMEOUT} and triggers failover. + */ + public Duration attemptTimeout() { + return attemptTimeout; + } + + public boolean isTimeoutEnabled() { + return attemptTimeout != null && !attemptTimeout.isZero() && !attemptTimeout.isNegative(); + } + + /** + * Whether an attempt that completes without emitting any response should be treated as a failure + * and trigger failover. This catches providers (Azure/Bedrock streaming) that historically + * completed silently on a non-2xx response. + */ + public boolean treatEmptyResponseAsFailure() { + return treatEmptyResponseAsFailure; + } + + /** Decides whether a given error should trigger failover to the next model in the chain. */ + public boolean shouldFailover(LlmError error) { + if (retryableCategories != null) { + return retryableCategories.contains(error.category()); + } + return error.isRetryable(); + } + + /** Builder for {@link FailoverConfig}. */ + public static final class Builder { + private Duration attemptTimeout = Duration.ZERO; + private boolean treatEmptyResponseAsFailure = true; + @Nullable private Set retryableCategories; + + private Builder() {} + + public Builder attemptTimeout(Duration attemptTimeout) { + this.attemptTimeout = attemptTimeout == null ? Duration.ZERO : attemptTimeout; + return this; + } + + public Builder attemptTimeoutMillis(long millis) { + this.attemptTimeout = Duration.ofMillis(Math.max(0, millis)); + return this; + } + + public Builder treatEmptyResponseAsFailure(boolean treatEmptyResponseAsFailure) { + this.treatEmptyResponseAsFailure = treatEmptyResponseAsFailure; + return this; + } + + /** + * Overrides which categories trigger failover. When unset, each category's default ({@link + * LlmFailureCategory#isRetryableByDefault()}) is used. + */ + public Builder retryableCategories(@Nullable Set retryableCategories) { + this.retryableCategories = retryableCategories; + return this; + } + + public FailoverConfig build() { + return new FailoverConfig(this); + } + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/FailoverLlm.java b/core/src/main/java/com/google/adk/models/failover/FailoverLlm.java new file mode 100644 index 000000000..3af99527c --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/FailoverLlm.java @@ -0,0 +1,318 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.common.base.Suppliers; +import io.reactivex.rxjava3.core.Flowable; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; + +/** + * A {@link BaseLlm} decorator that wraps a primary model plus an ordered list of fallback models + * and transparently fails over between them, in the same spirit as a Java fallback-API chain. + * + *

Key properties: + * + *

    + *
  • Zero happy-path latency. The primary model's {@code Flowable} is returned with only + * lightweight operators attached; fallbacks are attempted lazily and only on a classified, + * retryable failure or timeout. + *
  • Transparent to all agents. Because it implements {@link BaseLlm}, any {@code + * LlmAgent} (including those inside {@code SequentialAgent}, {@code ParallelAgent}, {@code + * LoopAgent} or custom orchestrators) gets failover simply by resolving to a {@code + * FailoverLlm} instance. + *
  • Uniform failure handling. Three failure shapes are normalized into one path: + * exceptions on the {@code Flowable}, in-band {@link LlmResponse#errorCode()} emissions, and + * silent empty completion (Azure/Bedrock). + *
+ * + *

Streaming caveat: if a failure surfaces after partial chunks were already emitted, downstream + * consumers will have observed those partial chunks before the fallback model's output. Failover is + * cleanest for non-streaming calls, which is the dominant path for the agents in this codebase. + */ +public class FailoverLlm extends BaseLlm { + + private final List> chain; + private final FailoverConfig config; + + /** Builds a failover chain from already-instantiated models. */ + public FailoverLlm(BaseLlm primary, BaseLlm... fallbacks) { + this(primary, Arrays.asList(fallbacks), FailoverConfig.defaults()); + } + + /** Builds a failover chain from already-instantiated models with explicit config. */ + public FailoverLlm(BaseLlm primary, List fallbacks, FailoverConfig config) { + this(toSuppliers(primary, fallbacks), config); + } + + /** + * Builds a failover chain from lazy suppliers. The first supplier is the primary; each supplier + * is memoized so the underlying model (and its HTTP client) is built at most once. + * + * @param chain ordered suppliers, index 0 = primary; must be non-empty. + * @param config failover tuning. + */ + public FailoverLlm(List> chain, FailoverConfig config) { + super(primaryModelName(chain)); + List> memoized = new ArrayList<>(chain.size()); + for (Supplier supplier : chain) { + memoized.add(Suppliers.memoize(supplier::get)::get); + } + this.chain = memoized; + this.config = config == null ? FailoverConfig.defaults() : config; + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return attempt(0, llmRequest, stream); + } + + private Flowable attempt(int index, LlmRequest llmRequest, boolean stream) { + final int attemptIndex = index; + final BaseLlm attemptLlm = resolve(attemptIndex); + + // IMPORTANT: Ensure the per-attempt request's model name matches the model being invoked. + // + // Many call sites set LlmRequest.model() (e.g. to the current model's name). Without rewriting + // this field, a cross-provider failover chain can accidentally call the fallback provider with + // the primary provider's model identifier (e.g. Azure deployment name leaking into Gemini), + // leading to confusing "model not found" errors on the fallback attempt. + final LlmRequest attemptRequest = + llmRequest.model().isPresent() + ? llmRequest.toBuilder().model(attemptLlm.model()).build() + : llmRequest; + + // Defer so the underlying call (and any eager work it does on subscription) is only triggered + // when this attempt is actually reached. + Flowable flow = + Flowable.defer(() -> attemptLlm.generateContent(attemptRequest, stream)); + + final String attemptModel = safeModelName(attemptIndex); + + // 1) Convert in-band error responses (errorCode set) into exceptions so they share the failover + // path with thrown errors. + flow = + flow.map( + response -> { + LlmErrorClassifier.classifyInBand(attemptModel, response) + .filter(config::shouldFailover) + .ifPresent( + error -> { + throw new InBandFailureException(error); + }); + return response; + }); + + // 2) Treat silent empty completion (e.g. Azure/Bedrock non-2xx) as a failure. + if (config.treatEmptyResponseAsFailure()) { + flow = + flow.switchIfEmpty( + Flowable.error( + () -> + new LlmHttpException( + -1, "Model returned an empty response: " + attemptModel, null))); + } + + // 3) Latency budget -> TimeoutException -> classified as TIMEOUT. + if (config.isTimeoutEnabled()) { + flow = flow.timeout(config.attemptTimeout().toMillis(), TimeUnit.MILLISECONDS); + } + + // 4) Record fallback success exactly once for non-primary attempts (attached before failover so + // it only fires for this model's own emissions). + if (attemptIndex > 0) { + AtomicBoolean recorded = new AtomicBoolean(false); + flow = + flow.doOnNext( + response -> { + if (recorded.compareAndSet(false, true)) { + LlmCallMetrics.recordFallbackSuccess(attemptModel, model(), attemptIndex); + } + }); + } + + // 5) Failover. + return flow.onErrorResumeNext( + throwable -> { + LlmError error = + throwable instanceof InBandFailureException + ? ((InBandFailureException) throwable).error() + : LlmErrorClassifier.classify(attemptModel, throwable); + + LlmCallMetrics.recordFailure(error, attemptIndex, model()); + + boolean hasNext = attemptIndex + 1 < chain.size(); + if (hasNext && config.shouldFailover(error)) { + return attempt(attemptIndex + 1, llmRequest, stream); + } + + if (!hasNext) { + LlmCallMetrics.recordExhausted(error, model(), chain.size()); + } + // Non-retryable, or chain exhausted: propagate the original cause. + return Flowable.error(unwrapForPropagation(throwable)); + }); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + // Live/bidi failover is out of scope; delegate to the primary model. + return resolve(0).connect(llmRequest); + } + + /** + * Returns the resolved primary model instance (attempt index 0). + * + *

This is primarily intended for application-level wiring code that needs to unwrap a + * previously-wrapped {@link FailoverLlm} without losing provider-specific identity (e.g. Azure + * deployment names that are not valid {@link com.google.adk.models.LlmRegistry} keys by + * themselves). + */ + public BaseLlm primary() { + return resolve(0); + } + + /** The ordered model names in the failover chain (primary first). For diagnostics. */ + public List chainModelNames() { + List names = new ArrayList<>(chain.size()); + for (int i = 0; i < chain.size(); i++) { + names.add(safeModelName(i)); + } + return names; + } + + private BaseLlm resolve(int index) { + return Objects.requireNonNull( + chain.get(index).get(), "Failover chain supplier returned null at index " + index); + } + + private String safeModelName(int index) { + try { + return resolve(index).model(); + } catch (RuntimeException e) { + return ""; + } + } + + private static Throwable unwrapForPropagation(Throwable throwable) { + if (throwable instanceof InBandFailureException) { + LlmError error = ((InBandFailureException) throwable).error(); + return error + .cause() + .orElseGet(() -> new IllegalStateException("LLM in-band error: " + error.message())); + } + return throwable; + } + + private static List> toSuppliers(BaseLlm primary, List fallbacks) { + List> suppliers = new ArrayList<>(); + suppliers.add(() -> primary); + if (fallbacks != null) { + for (BaseLlm fallback : fallbacks) { + suppliers.add(() -> fallback); + } + } + return suppliers; + } + + private static String primaryModelName(List> chain) { + if (chain == null || chain.isEmpty()) { + throw new IllegalArgumentException("FailoverLlm requires at least a primary model."); + } + try { + BaseLlm primary = chain.get(0).get(); + return primary == null ? "failover" : primary.model(); + } catch (RuntimeException e) { + // Primary not yet resolvable (lazy); name is informational only. + return "failover"; + } + } + + /** Carries a normalized in-band {@link LlmError} through the RxJava error channel. */ + private static final class InBandFailureException extends RuntimeException { + private static final long serialVersionUID = 1L; + private final transient LlmError error; + + InBandFailureException(LlmError error) { + super(error.message()); + this.error = error; + } + + LlmError error() { + return error; + } + } + + /** Fluent builder for {@link FailoverLlm}. */ + public static final class Builder { + private final List> chain = new ArrayList<>(); + private FailoverConfig config = FailoverConfig.defaults(); + + private Builder() {} + + public Builder primary(BaseLlm primary) { + Objects.requireNonNull(primary, "primary"); + add(() -> primary); + return this; + } + + public Builder primary(Supplier primary) { + add(primary); + return this; + } + + public Builder fallback(BaseLlm fallback) { + Objects.requireNonNull(fallback, "fallback"); + add(() -> fallback); + return this; + } + + public Builder fallback(Supplier fallback) { + add(fallback); + return this; + } + + public Builder config(FailoverConfig config) { + this.config = config == null ? FailoverConfig.defaults() : config; + return this; + } + + private void add(Supplier supplier) { + chain.add(Objects.requireNonNull(supplier, "model supplier")); + } + + public FailoverLlm build() { + if (chain.isEmpty()) { + throw new IllegalStateException("FailoverLlm.Builder requires at least a primary model."); + } + return new FailoverLlm(chain, config); + } + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/LlmCallMetrics.java b/core/src/main/java/com/google/adk/models/failover/LlmCallMetrics.java new file mode 100644 index 000000000..54d25528a --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/LlmCallMetrics.java @@ -0,0 +1,123 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Pluggable sink for LLM failover telemetry. The platform deliberately avoids a hard dependency on + * any specific metrics library; applications (e.g. rae) can register a {@link Recorder} that + * bridges to Micrometer, Prometheus, Datadog, etc. + * + *

The default {@link LoggingRecorder} emits structured SLF4J logs so that provider error codes + * are tracked out of the box even before an application wires its own metrics backend. + */ +public final class LlmCallMetrics { + + /** Receives failover lifecycle events. All methods have no-op defaults. */ + public interface Recorder { + /** Invoked for every failed attempt, whether or not a fallback follows. */ + default void onFailure(LlmError error, int attemptIndex, String primaryModel) {} + + /** Invoked when a fallback model succeeds after the primary (or an earlier fallback) failed. */ + default void onFallbackSuccess(String succeededModel, String primaryModel, int attemptIndex) {} + + /** Invoked when every model in the chain has failed and the error is about to propagate. */ + default void onExhausted(LlmError lastError, String primaryModel, int totalAttempts) {} + } + + private static volatile Recorder recorder = new LoggingRecorder(); + + private LlmCallMetrics() {} + + /** Replaces the active recorder. Pass an application-specific implementation at startup. */ + public static void setRecorder(Recorder newRecorder) { + recorder = newRecorder == null ? new LoggingRecorder() : newRecorder; + } + + /** Returns the active recorder. */ + public static Recorder recorder() { + return recorder; + } + + static void recordFailure(LlmError error, int attemptIndex, String primaryModel) { + try { + recorder.onFailure(error, attemptIndex, primaryModel); + } catch (RuntimeException e) { + LoggingRecorder.LOGGER.warn("LlmCallMetrics recorder threw on onFailure", e); + } + } + + static void recordFallbackSuccess(String succeededModel, String primaryModel, int attemptIndex) { + try { + recorder.onFallbackSuccess(succeededModel, primaryModel, attemptIndex); + } catch (RuntimeException e) { + LoggingRecorder.LOGGER.warn("LlmCallMetrics recorder threw on onFallbackSuccess", e); + } + } + + static void recordExhausted(LlmError lastError, String primaryModel, int totalAttempts) { + try { + recorder.onExhausted(lastError, primaryModel, totalAttempts); + } catch (RuntimeException e) { + LoggingRecorder.LOGGER.warn("LlmCallMetrics recorder threw on onExhausted", e); + } + } + + /** Default recorder that logs structured failover telemetry via SLF4J. */ + public static final class LoggingRecorder implements Recorder { + static final Logger LOGGER = LoggerFactory.getLogger("com.google.adk.models.failover"); + + @Override + public void onFailure(LlmError error, int attemptIndex, String primaryModel) { + LOGGER.warn( + "LLM attempt failed [primary={}, attemptModel={}, attemptIndex={}, provider={}," + + " httpStatus={}, providerCode={}, category={}, retryable={}, message={}]", + primaryModel, + error.modelName(), + attemptIndex, + error.provider().orElse("unknown"), + error.httpStatus(), + error.providerCode().orElse("n/a"), + error.category(), + error.isRetryable(), + error.message()); + } + + @Override + public void onFallbackSuccess(String succeededModel, String primaryModel, int attemptIndex) { + LOGGER.info( + "LLM fallback succeeded [primary={}, succeededModel={}, attemptIndex={}]", + primaryModel, + succeededModel, + attemptIndex); + } + + @Override + public void onExhausted(LlmError lastError, String primaryModel, int totalAttempts) { + LOGGER.error( + "LLM failover exhausted [primary={}, totalAttempts={}, lastCategory={}, lastHttpStatus={}," + + " lastMessage={}]", + primaryModel, + totalAttempts, + lastError.category(), + lastError.httpStatus(), + lastError.message()); + } + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/LlmError.java b/core/src/main/java/com/google/adk/models/failover/LlmError.java new file mode 100644 index 000000000..59381c67c --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/LlmError.java @@ -0,0 +1,166 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +import java.util.Optional; +import org.jspecify.annotations.Nullable; + +/** + * Normalized, provider-agnostic description of a single failed LLM attempt. + * + *

This is the unit recorded by {@link LlmCallMetrics} and inspected by {@link FailoverLlm} to + * decide whether to fail over. It captures the HTTP status code where available, a provider-native + * code/status string, the normalized {@link LlmFailureCategory}, a human readable message, and the + * underlying cause. + */ +public final class LlmError { + + private final String modelName; + @Nullable private final String provider; + private final int httpStatus; + @Nullable private final String providerCode; + private final LlmFailureCategory category; + private final String message; + @Nullable private final Throwable cause; + + private LlmError( + String modelName, + @Nullable String provider, + int httpStatus, + @Nullable String providerCode, + LlmFailureCategory category, + String message, + @Nullable Throwable cause) { + this.modelName = modelName; + this.provider = provider; + this.httpStatus = httpStatus; + this.providerCode = providerCode; + this.category = category; + this.message = message; + this.cause = cause; + } + + public static Builder builder(String modelName) { + return new Builder(modelName); + } + + /** The model name this failure is attributed to. */ + public String modelName() { + return modelName; + } + + /** Best-effort provider identifier (e.g. "gemini", "azure", "bedrock"). */ + public Optional provider() { + return Optional.ofNullable(provider); + } + + /** HTTP status code, or {@code -1} if not applicable / unknown. */ + public int httpStatus() { + return httpStatus; + } + + /** Provider-native error code or status string (e.g. "RESOURCE_EXHAUSTED"). */ + public Optional providerCode() { + return Optional.ofNullable(providerCode); + } + + /** Normalized failure category. */ + public LlmFailureCategory category() { + return category; + } + + /** Human readable failure message. */ + public String message() { + return message; + } + + /** Underlying throwable, if this error originated from an exception. */ + public Optional cause() { + return Optional.ofNullable(cause); + } + + /** Whether failing over to another model is sensible for this error, per its category. */ + public boolean isRetryable() { + return category.isRetryableByDefault(); + } + + @Override + public String toString() { + return "LlmError{model=" + + modelName + + ", provider=" + + provider + + ", httpStatus=" + + httpStatus + + ", providerCode=" + + providerCode + + ", category=" + + category + + ", message=" + + message + + '}'; + } + + /** Mutable builder for {@link LlmError}. */ + public static final class Builder { + private final String modelName; + @Nullable private String provider; + private int httpStatus = -1; + @Nullable private String providerCode; + private LlmFailureCategory category = LlmFailureCategory.UNKNOWN; + private String message = ""; + @Nullable private Throwable cause; + + private Builder(String modelName) { + this.modelName = modelName; + } + + public Builder provider(@Nullable String provider) { + this.provider = provider; + return this; + } + + public Builder httpStatus(int httpStatus) { + this.httpStatus = httpStatus; + return this; + } + + public Builder providerCode(@Nullable String providerCode) { + this.providerCode = providerCode; + return this; + } + + public Builder category(LlmFailureCategory category) { + this.category = category; + return this; + } + + public Builder message(@Nullable String message) { + this.message = message == null ? "" : message; + return this; + } + + public Builder cause(@Nullable Throwable cause) { + this.cause = cause; + return this; + } + + public LlmError build() { + return new LlmError(modelName, provider, httpStatus, providerCode, category, message, cause); + } + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/LlmErrorClassifier.java b/core/src/main/java/com/google/adk/models/failover/LlmErrorClassifier.java new file mode 100644 index 000000000..67caa1cc5 --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/LlmErrorClassifier.java @@ -0,0 +1,214 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +import com.google.adk.models.LlmResponse; +import com.google.genai.errors.ApiException; +import java.io.IOException; +import java.util.Locale; +import java.util.Optional; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Normalizes the heterogeneous ways adk-java model implementations report failures into a single + * {@link LlmError}. + * + *

Failures arrive in three different shapes across providers: + * + *

    + *
  • Thrown exceptions on the {@code Flowable} (genai {@link ApiException} with an HTTP code, + * {@link LlmHttpException} from the Azure/Bedrock transports, {@link IOException}, generic + * {@link IllegalStateException} with the status embedded in the message, RxJava {@link + * TimeoutException}). + *
  • In-band "successful" emissions where {@link LlmResponse#errorCode()} is set (the genai path + * emits these instead of throwing). + *
  • Silent empty completion (handled by {@link FailoverLlm}, not here). + *
+ */ +public final class LlmErrorClassifier { + + // Matches "status=429", "status: 503", "code 500", "HTTP 429", "(429)" etc. + private static final Pattern STATUS_PATTERN = + Pattern.compile("(?:status|code|http)[^0-9]{0,3}(\\d{3})", Pattern.CASE_INSENSITIVE); + private static final Pattern BARE_STATUS_PATTERN = Pattern.compile("\\b([45]\\d{2})\\b"); + + private LlmErrorClassifier() {} + + /** Classifies a throwable raised by a model call into a normalized {@link LlmError}. */ + public static LlmError classify(String modelName, Throwable throwable) { + Throwable t = unwrap(throwable); + LlmError.Builder builder = + LlmError.builder(modelName).cause(t).message(messageOf(t)).provider(providerOf(modelName)); + + if (t instanceof TimeoutException) { + return builder.category(LlmFailureCategory.TIMEOUT).build(); + } + + if (t instanceof LlmHttpException) { + int status = ((LlmHttpException) t).statusCode(); + return builder.httpStatus(status).category(categoryForStatus(status)).build(); + } + + if (t instanceof ApiException) { + ApiException api = (ApiException) t; + return builder + .httpStatus(api.code()) + .providerCode(api.status()) + .category(categoryForStatus(api.code())) + .build(); + } + + // genai wraps connectivity issues in GenAiIOException; detect by type hierarchy / name. + if (t instanceof IOException || isGenAiIoException(t)) { + Optional status = extractStatus(messageOf(t)); + if (status.isPresent()) { + return builder.httpStatus(status.get()).category(categoryForStatus(status.get())).build(); + } + return builder.category(LlmFailureCategory.NETWORK).build(); + } + + // Azure/Bedrock legacy paths throw IllegalStateException with the status in the message. + Optional status = extractStatus(messageOf(t)); + if (status.isPresent()) { + return builder.httpStatus(status.get()).category(categoryForStatus(status.get())).build(); + } + + return builder.category(LlmFailureCategory.UNKNOWN).build(); + } + + /** + * Classifies an in-band error carried on a "successful" {@link LlmResponse} (i.e. {@link + * LlmResponse#errorCode()} is present). Returns empty if the response is not an error. + */ + public static Optional classifyInBand(String modelName, LlmResponse response) { + if (response.errorCode().isEmpty()) { + return Optional.empty(); + } + String code = response.errorCode().get().toString(); + String msg = response.errorMessage().orElse(code); + LlmFailureCategory category = categoryForFinishReason(code); + return Optional.of( + LlmError.builder(modelName) + .provider(providerOf(modelName)) + .providerCode(code) + .category(category) + .message(msg) + .build()); + } + + /** Maps an HTTP status code to a normalized failure category. */ + public static LlmFailureCategory categoryForStatus(int status) { + if (status == 429) { + return LlmFailureCategory.RATE_LIMIT; + } + if (status == 503) { + return LlmFailureCategory.UNAVAILABLE; + } + if (status == 408) { + return LlmFailureCategory.TIMEOUT; + } + if (status == 401 || status == 403) { + return LlmFailureCategory.AUTH; + } + if (status >= 500 && status < 600) { + return LlmFailureCategory.SERVER_ERROR; + } + if (status >= 400 && status < 500) { + return LlmFailureCategory.BAD_REQUEST; + } + return LlmFailureCategory.UNKNOWN; + } + + private static LlmFailureCategory categoryForFinishReason(String reason) { + String upper = reason.toUpperCase(Locale.ROOT); + if (upper.contains("SAFETY") + || upper.contains("BLOCK") + || upper.contains("RECITATION") + || upper.contains("PROHIBITED") + || upper.contains("SPII")) { + return LlmFailureCategory.CONTENT_BLOCKED; + } + return LlmFailureCategory.UNKNOWN; + } + + private static Optional extractStatus(String message) { + if (message == null || message.isEmpty()) { + return Optional.empty(); + } + Matcher m = STATUS_PATTERN.matcher(message); + if (m.find()) { + return Optional.of(Integer.parseInt(m.group(1))); + } + Matcher bare = BARE_STATUS_PATTERN.matcher(message); + if (bare.find()) { + return Optional.of(Integer.parseInt(bare.group(1))); + } + return Optional.empty(); + } + + private static Throwable unwrap(Throwable t) { + Throwable current = t; + while ((current instanceof CompletionException || current instanceof ExecutionException) + && current.getCause() != null + && current.getCause() != current) { + current = current.getCause(); + } + return current; + } + + private static boolean isGenAiIoException(Throwable t) { + return t.getClass().getName().equals("com.google.genai.errors.GenAiIOException"); + } + + private static String messageOf(Throwable t) { + String msg = t.getMessage(); + return msg == null ? t.getClass().getSimpleName() : msg; + } + + private static String providerOf(String modelName) { + if (modelName == null) { + return null; + } + String lower = modelName.toLowerCase(Locale.ROOT); + if (lower.startsWith("gemini") || lower.startsWith("gemma")) { + return "google"; + } + if (lower.startsWith("azure") || lower.contains("realtime")) { + return "azure"; + } + if (lower.startsWith("bedrock") || lower.contains("anthropic") || lower.contains("claude")) { + return "bedrock"; + } + if (lower.startsWith("redbusadg")) { + return "redbusadg"; + } + if (lower.startsWith("ollama")) { + return "ollama"; + } + if (lower.startsWith("apigee")) { + return "apigee"; + } + if (lower.startsWith("gpt-oss")) { + return "gpt-oss"; + } + return null; + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/LlmFailureCategory.java b/core/src/main/java/com/google/adk/models/failover/LlmFailureCategory.java new file mode 100644 index 000000000..564beeddb --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/LlmFailureCategory.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +/** + * Normalized classification of an LLM call failure across heterogeneous providers (Gemini/genai, + * Azure, Bedrock, Claude, Ollama, etc.). + * + *

The {@code retryableByDefault} flag expresses whether a failure of this category is, in + * general, worth failing over to another model. Transient/availability failures are retryable; + * deterministic client-side failures (bad request, auth, content blocked) are not, because retrying + * the same request against a different model would almost certainly fail the same way. + */ +public enum LlmFailureCategory { + /** HTTP 429 / quota exhausted. Worth failing over to a different model or provider. */ + RATE_LIMIT(true), + + /** HTTP 503 / provider explicitly unavailable or overloaded. */ + UNAVAILABLE(true), + + /** HTTP 5xx (500, 502, 504, ...) server-side error. */ + SERVER_ERROR(true), + + /** Call exceeded the configured latency budget (RxJava {@code TimeoutException}). */ + TIMEOUT(true), + + /** Connection reset, DNS, socket errors, or other transport-level IO failures. */ + NETWORK(true), + + /** HTTP 401 / 403 authentication or authorization failure. */ + AUTH(true), + + /** HTTP 400 / 404 / 422 and similar deterministic client errors. */ + BAD_REQUEST(true), + + /** Prompt or response blocked by safety / recitation / content filters. */ + CONTENT_BLOCKED(false), + + /** Failure that could not be classified into any of the above. */ + UNKNOWN(true); + + private final boolean retryableByDefault; + + LlmFailureCategory(boolean retryableByDefault) { + this.retryableByDefault = retryableByDefault; + } + + /** + * Whether failing over to another model is sensible for this category by default. Callers may + * override this decision via configuration. + */ + public boolean isRetryableByDefault() { + return retryableByDefault; + } +} diff --git a/core/src/main/java/com/google/adk/models/failover/LlmHttpException.java b/core/src/main/java/com/google/adk/models/failover/LlmHttpException.java new file mode 100644 index 000000000..4461a6edf --- /dev/null +++ b/core/src/main/java/com/google/adk/models/failover/LlmHttpException.java @@ -0,0 +1,54 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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 com.google.adk.models.failover; + +/** + * Typed exception carrying an HTTP status code for providers that otherwise swallow or stringify + * transport failures (e.g. the Azure and Bedrock REST transports). + * + *

Surfacing the status code as a first-class field lets {@link LlmErrorClassifier} reliably bin + * to a {@link LlmFailureCategory} so that {@link FailoverLlm} can decide whether to fail over. + */ +public class LlmHttpException extends RuntimeException { + + private static final long serialVersionUID = 1L; + + private final int statusCode; + private final String responseBody; + + public LlmHttpException(int statusCode, String message, String responseBody) { + super("HTTP " + statusCode + (message != null ? ": " + message : "")); + this.statusCode = statusCode; + this.responseBody = responseBody; + } + + public LlmHttpException(int statusCode, String message, String responseBody, Throwable cause) { + super("HTTP " + statusCode + (message != null ? ": " + message : ""), cause); + this.statusCode = statusCode; + this.responseBody = responseBody; + } + + /** The HTTP status code returned by the provider, or {@code -1} if unknown. */ + public int statusCode() { + return statusCode; + } + + /** The raw response body, if captured. May be {@code null}. */ + public String responseBody() { + return responseBody; + } +}