Skip to content

refactor:Update BasePlugin to accept ToolContext.Builder for callbacks instead of immutable ToolContext #645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+154 −109
Open

refactor:Update BasePlugin to accept ToolContext.Builder for callbacks instead of immutable ToolContext #645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions core/src/main/java/com/google/adk/agents/Callbacks.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,14 +129,14 @@ public interface BeforeToolCallback extends BeforeToolCallbackBase {
* @param invocationContext Invocation context.
* @param baseTool Tool instance.
* @param input Tool input arguments.
* @param toolContext Tool context.
* @param toolContext Tool context builder.
* @return override result, or empty to continue.
*/
Maybe<Map<String, Object>> call(
InvocationContext invocationContext,
BaseTool baseTool,
Map<String, Object> input,
ToolContext toolContext);
ToolContext.Builder toolContext);
}

/**
Expand All @@ -149,7 +149,7 @@ Optional<Map<String, Object>> call(
InvocationContext invocationContext,
BaseTool baseTool,
Map<String, Object> input,
ToolContext toolContext);
ToolContext.Builder toolContext);
}

interface AfterToolCallbackBase {}
Expand All @@ -162,15 +162,15 @@ public interface AfterToolCallback extends AfterToolCallbackBase {
* @param invocationContext Invocation context.
* @param baseTool Tool instance.
* @param input Tool input arguments.
* @param toolContext Tool context.
* @param toolContext Tool context builder.
* @param response Raw tool response.
* @return processed result, or empty to keep original.
*/
Maybe<Map<String, Object>> call(
InvocationContext invocationContext,
BaseTool baseTool,
Map<String, Object> input,
ToolContext toolContext,
ToolContext.Builder toolContext,
Object response);
}

Expand All @@ -184,7 +184,7 @@ Optional<Map<String, Object>> call(
InvocationContext invocationContext,
BaseTool baseTool,
Map<String, Object> input,
ToolContext toolContext,
ToolContext.Builder toolContext,
Object response);
}

Expand Down
9 changes: 6 additions & 3 deletions core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,15 +231,18 @@ private Flowable<LlmResponse> callLlm(
.runOnModelErrorCallback(
new CallbackContext(
context, eventForCallbackUsage.actions()),
llmRequest,
llmRequestBuilder,
exception)
.switchIfEmpty(Single.error(exception))
.toFlowable())
.doOnNext(
llmResp -> {
try (Scope innerScope = llmCallSpan.makeCurrent()) {
Telemetry.traceCallLlm(
context, eventForCallbackUsage.id(), llmRequest, llmResp);
context,
eventForCallbackUsage.id(),
llmRequestBuilder.build(),
llmResp);
}
})
.doOnError(
Expand Down Expand Up @@ -269,7 +272,7 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());

Maybe<LlmResponse> pluginResult =
context.pluginManager().runBeforeModelCallback(callbackContext, llmRequestBuilder.build());
context.pluginManager().runBeforeModelCallback(callbackContext, llmRequestBuilder);

LlmAgent agent = (LlmAgent) context.agent();

Expand Down
42 changes: 23 additions & 19 deletions core/src/main/java/com/google/adk/flows/llmflows/Functions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,17 +147,17 @@ public static Maybe<Event> handleFunctionCalls(
Function<FunctionCall, Maybe<Event>> functionCallMapper =
functionCall -> {
BaseTool tool = tools.get(functionCall.name().get());
ToolContext toolContext =
ToolContext.Builder toolContextBuilder =
ToolContext.builder(invocationContext)
.functionCallId(functionCall.id().orElse(""))
.toolConfirmation(toolConfirmations.get(functionCall.id().orElse(null)))
.build();
.toolConfirmation(toolConfirmations.get(functionCall.id().orElse(null)));

Map<String, Object> functionArgs = functionCall.args().orElse(ImmutableMap.of());

Maybe<Map<String, Object>> maybeFunctionResult =
maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
.switchIfEmpty(Maybe.defer(() -> callTool(tool, functionArgs, toolContext)));
maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContextBuilder)
.switchIfEmpty(
Maybe.defer(() -> callTool(tool, functionArgs, toolContextBuilder.build())));

return maybeFunctionResult
.map(Optional::of)
Expand All @@ -166,7 +166,7 @@ public static Maybe<Event> handleFunctionCalls(
t ->
invocationContext
.pluginManager()
.runOnToolErrorCallback(tool, functionArgs, toolContext, t)
.runOnToolErrorCallback(tool, functionArgs, toolContextBuilder, t)
.map(Optional::of)
.switchIfEmpty(Single.error(t)))
.flatMapMaybe(
Expand All @@ -178,7 +178,7 @@ public static Maybe<Event> handleFunctionCalls(
invocationContext,
tool,
functionArgs,
toolContext,
toolContextBuilder,
initialFunctionResult);

return afterToolResultMaybe
Expand All @@ -193,7 +193,10 @@ public static Maybe<Event> handleFunctionCalls(
}
Event functionResponseEvent =
buildResponseEvent(
tool, finalFunctionResult, toolContext, invocationContext);
tool,
finalFunctionResult,
toolContextBuilder.build(),
invocationContext);
return Maybe.just(functionResponseEvent);
});
});
Expand Down Expand Up @@ -252,21 +255,19 @@ public static Maybe<Event> handleFunctionCallsLive(
Function<FunctionCall, Maybe<Event>> functionCallMapper =
functionCall -> {
BaseTool tool = tools.get(functionCall.name().get());
ToolContext toolContext =
ToolContext.builder(invocationContext)
.functionCallId(functionCall.id().orElse(""))
.build();
ToolContext.Builder toolContextBuilder =
ToolContext.builder(invocationContext).functionCallId(functionCall.id().orElse(""));
Map<String, Object> functionArgs = functionCall.args().orElse(new HashMap<>());

Maybe<Map<String, Object>> maybeFunctionResult =
maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContextBuilder)
.switchIfEmpty(
Maybe.defer(
() ->
processFunctionLive(
invocationContext,
tool,
toolContext,
toolContextBuilder.build(),
functionCall,
functionArgs)));

Expand All @@ -277,7 +278,7 @@ public static Maybe<Event> handleFunctionCallsLive(
t ->
invocationContext
.pluginManager()
.runOnToolErrorCallback(tool, functionArgs, toolContext, t)
.runOnToolErrorCallback(tool, functionArgs, toolContextBuilder, t)
.map(Optional::ofNullable)
.switchIfEmpty(Single.error(t)))
.flatMapMaybe(
Expand All @@ -289,7 +290,7 @@ public static Maybe<Event> handleFunctionCallsLive(
invocationContext,
tool,
functionArgs,
toolContext,
toolContextBuilder,
initialFunctionResult);

return afterToolResultMaybe
Expand All @@ -304,7 +305,10 @@ public static Maybe<Event> handleFunctionCallsLive(
}
Event functionResponseEvent =
buildResponseEvent(
tool, finalFunctionResult, toolContext, invocationContext);
tool,
finalFunctionResult,
toolContextBuilder.build(),
invocationContext);
return Maybe.just(functionResponseEvent);
});
});
Expand Down Expand Up @@ -466,7 +470,7 @@ private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
InvocationContext invocationContext,
BaseTool tool,
Map<String, Object> functionArgs,
ToolContext toolContext) {
ToolContext.Builder toolContext) {
if (invocationContext.agent() instanceof LlmAgent) {
LlmAgent agent = (LlmAgent) invocationContext.agent();

Expand Down Expand Up @@ -497,7 +501,7 @@ private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
InvocationContext invocationContext,
BaseTool tool,
Map<String, Object> functionArgs,
ToolContext toolContext,
ToolContext.Builder toolContext,
Map<String, Object> functionResult) {
if (invocationContext.agent() instanceof LlmAgent) {
LlmAgent agent = (LlmAgent) invocationContext.agent();
Expand Down
25 changes: 15 additions & 10 deletions core/src/main/java/com/google/adk/plugins/BasePlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,11 +122,12 @@ public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callba
* Callback executed before a request is sent to the model.
*
* @param callbackContext The context for the current agent call.
* @param llmRequest The prepared request object to be sent to the model.
* @param llmRequest The mutable request builder, allowing modification of the request before it
* is sent to the model.
* @return An optional LlmResponse to trigger an early exit. Returning Empty to proceed normally.
*/
public Maybe<LlmResponse> beforeModelCallback(
CallbackContext callbackContext, LlmRequest llmRequest) {
CallbackContext callbackContext, LlmRequest.Builder llmRequest) {
return Maybe.empty();
}

Expand All @@ -147,13 +148,13 @@ public Maybe<LlmResponse> afterModelCallback(
* Callback executed when a model call encounters an error.
*
* @param callbackContext The context for the current agent call.
* @param llmRequest The request that was sent to the model.
* @param llmRequest The mutable request builder for the request that failed.
* @param error The exception that was raised.
* @return An optional LlmResponse to use instead of propagating the error. Returning Empty to
* allow the original error to be raised.
*/
public Maybe<LlmResponse> onModelErrorCallback(
CallbackContext callbackContext, LlmRequest llmRequest, Throwable error) {
CallbackContext callbackContext, LlmRequest.Builder llmRequest, Throwable error) {
return Maybe.empty();
}

Expand All @@ -162,12 +163,13 @@ public Maybe<LlmResponse> onModelErrorCallback(
*
* @param tool The tool instance that is about to be executed.
* @param toolArgs The dictionary of arguments to be used for invoking the tool.
* @param toolContext The context specific to the tool execution.
* @param toolContext The mutable tool context builder, allowing modification of the context
* before tool execution.
* @return An optional Map to stop the tool execution and return this response immediately.
* Returning Empty to proceed normally.
*/
public Maybe<Map<String, Object>> beforeToolCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext) {
BaseTool tool, Map<String, Object> toolArgs, ToolContext.Builder toolContext) {
return Maybe.empty();
}

Expand All @@ -176,15 +178,15 @@ public Maybe<Map<String, Object>> beforeToolCallback(
*
* @param tool The tool instance that has just been executed.
* @param toolArgs The original arguments that were passed to the tool.
* @param toolContext The context specific to the tool execution.
* @param toolContext The mutable tool context builder used for tool execution.
* @param result The dictionary returned by the tool invocation.
* @return An optional Map to replace the original result from the tool. Returning Empty to use
* the original result.
*/
public Maybe<Map<String, Object>> afterToolCallback(
BaseTool tool,
Map<String, Object> toolArgs,
ToolContext toolContext,
ToolContext.Builder toolContext,
Map<String, Object> result) {
return Maybe.empty();
}
Expand All @@ -194,13 +196,16 @@ public Maybe<Map<String, Object>> afterToolCallback(
*
* @param tool The tool instance that encountered an error.
* @param toolArgs The arguments that were passed to the tool.
* @param toolContext The context specific to the tool execution.
* @param toolContext The mutable tool context builder for the tool call that failed.
* @param error The exception that was raised during tool execution.
* @return An optional Map to be used as the tool response instead of propagating the error.
* Returning Empty to allow the original error to be raised.
*/
public Maybe<Map<String, Object>> onToolErrorCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext, Throwable error) {
BaseTool tool,
Map<String, Object> toolArgs,
ToolContext.Builder toolContext,
Throwable error) {
return Maybe.empty();
}
}
37 changes: 22 additions & 15 deletions core/src/main/java/com/google/adk/plugins/LoggingPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,14 +151,15 @@ public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callba

@Override
public Maybe<LlmResponse> beforeModelCallback(
CallbackContext callbackContext, LlmRequest llmRequest) {
CallbackContext callbackContext, LlmRequest.Builder llmRequest) {
return Maybe.fromAction(
() -> {
LlmRequest request = llmRequest.build();
log("🧠 LLM REQUEST");
log(" Model: " + llmRequest.model().orElse("default"));
log(" Model: " + request.model().orElse("default"));
log(" Agent: " + callbackContext.agentName());

llmRequest
request
.getFirstSystemInstruction()
.ifPresent(
sysInstruction -> {
Expand All @@ -170,8 +171,8 @@ public Maybe<LlmResponse> beforeModelCallback(
log(" System Instruction: '" + truncatedInstruction + "'");
});

if (!llmRequest.tools().isEmpty()) {
String toolNames = String.join(", ", llmRequest.tools().keySet());
if (!request.tools().isEmpty()) {
String toolNames = String.join(", ", request.tools().keySet());
log(" Available Tools: [" + toolNames + "]");
}
});
Expand Down Expand Up @@ -211,7 +212,7 @@ public Maybe<LlmResponse> afterModelCallback(

@Override
public Maybe<LlmResponse> onModelErrorCallback(
CallbackContext callbackContext, LlmRequest llmRequest, Throwable error) {
CallbackContext callbackContext, LlmRequest.Builder llmRequest, Throwable error) {
return Maybe.fromAction(
() -> {
log("🧠 LLM ERROR");
Expand All @@ -223,13 +224,14 @@ public Maybe<LlmResponse> onModelErrorCallback(

@Override
public Maybe<Map<String, Object>> beforeToolCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext) {
BaseTool tool, Map<String, Object> toolArgs, ToolContext.Builder toolContext) {
return Maybe.fromAction(
() -> {
ToolContext tc = toolContext.build();
log("🔧 TOOL STARTING");
log(" Tool Name: " + tool.name());
log(" Agent: " + toolContext.agentName());
toolContext.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Agent: " + tc.agentName());
tc.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Arguments: " + formatArgs(toolArgs));
});
}
Expand All @@ -238,27 +240,32 @@ public Maybe<Map<String, Object>> beforeToolCallback(
public Maybe<Map<String, Object>> afterToolCallback(
BaseTool tool,
Map<String, Object> toolArgs,
ToolContext toolContext,
ToolContext.Builder toolContext,
Map<String, Object> result) {
return Maybe.fromAction(
() -> {
ToolContext tc = toolContext.build();
log("🔧 TOOL COMPLETED");
log(" Tool Name: " + tool.name());
log(" Agent: " + toolContext.agentName());
toolContext.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Agent: " + tc.agentName());
tc.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Result: " + formatArgs(result));
});
}

@Override
public Maybe<Map<String, Object>> onToolErrorCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext, Throwable error) {
BaseTool tool,
Map<String, Object> toolArgs,
ToolContext.Builder toolContext,
Throwable error) {
return Maybe.fromAction(
() -> {
ToolContext tc = toolContext.build();
log("🔧 TOOL ERROR");
log(" Tool Name: " + tool.name());
log(" Agent: " + toolContext.agentName());
toolContext.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Agent: " + tc.agentName());
tc.functionCallId().ifPresent(id -> log(" Function Call ID: " + id));
log(" Arguments: " + formatArgs(toolArgs));
log(" Error: " + error.getMessage());
logger.error("[{}] Tool Error", name, error);
Expand Down
Loading
Loading