fix!: Allow beforeModelCallback to modify the LLM request

PiperOrigin-RevId: 804942304

Author: google-genai-bot

core/src/main/java/com/google/adk/agents/Callbacks.java

@@ -36,10 +36,10 @@ public interface BeforeModelCallback extends BeforeModelCallbackBase {
      * Async callback before LLM invocation.
      *
      * @param callbackContext Callback context.
-     * @param llmRequest LLM request.
+     * @param llmRequestBuilder LLM request builder.
      * @return response override, or empty to continue.
      */
-    Maybe<LlmResponse> call(CallbackContext callbackContext, LlmRequest llmRequest);
+    Maybe<LlmResponse> call(CallbackContext callbackContext, LlmRequest.Builder llmRequestBuilder);
   }
 
   /**
@@ -48,7 +48,8 @@ public interface BeforeModelCallback extends BeforeModelCallbackBase {
    */
   @FunctionalInterface
   public interface BeforeModelCallbackSync extends BeforeModelCallbackBase {
-    Optional<LlmResponse> call(CallbackContext callbackContext, LlmRequest llmRequest);
+    Optional<LlmResponse> call(
+        CallbackContext callbackContext, LlmRequest.Builder llmRequestBuilder);
   }
 
   interface AfterModelCallbackBase {}

core/src/main/java/com/google/adk/agents/LlmAgent.java

@@ -360,9 +360,10 @@ public Builder beforeModelCallback(List<BeforeModelCallbackBase> beforeModelCall
           } else if (callback instanceof BeforeModelCallbackSync beforeModelCallbackSyncInstance) {
             builder.add(
                 (BeforeModelCallback)
-                    (callbackContext, llmRequest) ->
+                    (callbackContext, llmRequestBuilder) ->
                         Maybe.fromOptional(
-                            beforeModelCallbackSyncInstance.call(callbackContext, llmRequest)));
+                            beforeModelCallbackSyncInstance.call(
+                                callbackContext, llmRequestBuilder)));
           } else {
             logger.warn(
                 "Invalid beforeModelCallback callback type: %s. Ignoring this callback.",
@@ -379,8 +380,9 @@ public Builder beforeModelCallback(List<BeforeModelCallbackBase> beforeModelCall
     public Builder beforeModelCallbackSync(BeforeModelCallbackSync beforeModelCallbackSync) {
       this.beforeModelCallback =
           ImmutableList.of(
-              (callbackContext, llmRequest) ->
-                  Maybe.fromOptional(beforeModelCallbackSync.call(callbackContext, llmRequest)));
+              (callbackContext, llmRequestBuilder) ->
+                  Maybe.fromOptional(
+                      beforeModelCallbackSync.call(callbackContext, llmRequestBuilder)));
       return this;
     }
 

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -209,7 +209,9 @@ private Flowable<LlmResponse> callLlm(
       InvocationContext context, LlmRequest llmRequest, Event eventForCallbackUsage) {
     LlmAgent agent = (LlmAgent) context.agent();
 
-    return handleBeforeModelCallback(context, llmRequest, eventForCallbackUsage)
+    LlmRequest.Builder llmRequestBuilder = llmRequest.toBuilder();
+
+    return handleBeforeModelCallback(context, llmRequestBuilder, eventForCallbackUsage)
         .flatMapPublisher(
             beforeResponse -> {
               if (beforeResponse.isPresent()) {
@@ -226,7 +228,7 @@ private Flowable<LlmResponse> callLlm(
 
                         try (Scope scope = llmCallSpan.makeCurrent()) {
                           return llm.generateContent(
-                                  llmRequest,
+                                  llmRequestBuilder.build(),
                                   context.runConfig().streamingMode() == StreamingMode.SSE)
                               .doOnNext(
                                   llmResp -> {
@@ -257,7 +259,7 @@ private Flowable<LlmResponse> callLlm(
    * @return A {@link Single} with the callback result or {@link Optional#empty()}.
    */
   private Single<Optional<LlmResponse>> handleBeforeModelCallback(
-      InvocationContext context, LlmRequest llmRequest, Event modelResponseEvent) {
+      InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) {
     LlmAgent agent = (LlmAgent) context.agent();
 
     Optional<List<BeforeModelCallback>> callbacksOpt = agent.beforeModelCallback();
@@ -274,7 +276,7 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
               CallbackContext callbackContext =
                   new CallbackContext(context, callbackEvent.actions());
               return callback
-                  .call(callbackContext, llmRequest)
+                  .call(callbackContext, llmRequestBuilder)
                   .map(Optional::of)
                   .defaultIfEmpty(Optional.empty());
             })

core/src/test/java/com/google/adk/agents/CallbacksTest.java

@@ -396,6 +396,27 @@ public void testRun_withBeforeModelCallback_returnsResponseFromCallback() {
     assertThat(getOnlyElement(events).content()).hasValue(callbackContent);
   }
 
+  @Test
+  public void testRun_withBeforeModelCallback_usesModifiedRequestFromCallback() {
+    TestLlm testLlm = createTestLlm(createLlmResponse(Content.builder().build()));
+    LlmAgent agent =
+        createTestAgentBuilder(testLlm)
+            .beforeModelCallback(
+                (context, requestBuilder) -> {
+                  requestBuilder.contents(
+                      ImmutableList.of(Content.fromParts(Part.fromText("Modified request"))));
+                  return Maybe.empty();
+                })
+            .build();
+    InvocationContext invocationContext = createInvocationContext(agent);
+
+    List<Event> unused = agent.runAsync(invocationContext).toList().blockingGet();
+
+    assertThat(testLlm.getRequests()).hasSize(1);
+    assertThat(testLlm.getRequests().get(0).contents())
+        .containsExactly(Content.fromParts(Part.fromText("Modified request")));
+  }
+
   @Test
   public void testRun_withAfterModelCallback_returnsResponseFromCallback() {
     Part textPartFromModel = Part.fromText("Real LLM response");