Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
10 changes: 5 additions & 5 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.events.Event;
import com.google.adk.plugins.PluginManager;
import com.google.adk.plugins.Plugin;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.errorprone.annotations.DoNotCall;
Expand Down Expand Up @@ -255,9 +255,9 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
* @return callback functions.
*/
private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacksToFunctions(
PluginManager pluginManager, List<? extends BeforeAgentCallback> callbacks) {
Plugin pluginManager, List<? extends BeforeAgentCallback> callbacks) {
return Stream.concat(
Stream.of(ctx -> pluginManager.runBeforeAgentCallback(this, ctx)),
Stream.of(ctx -> pluginManager.beforeAgentCallback(this, ctx)),
callbacks.stream()
.map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
.collect(toImmutableList());
Expand All @@ -270,9 +270,9 @@ private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacks
* @return callback functions.
*/
private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksToFunctions(
PluginManager pluginManager, List<? extends AfterAgentCallback> callbacks) {
Plugin pluginManager, List<? extends AfterAgentCallback> callbacks) {
return Stream.concat(
Stream.of(ctx -> pluginManager.runAfterAgentCallback(this, ctx)),
Stream.of(ctx -> pluginManager.afterAgentCallback(this, ctx)),
callbacks.stream()
.map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
.collect(toImmutableList());
Expand Down
11 changes: 6 additions & 5 deletions core/src/main/java/com/google/adk/agents/InvocationContext.java
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.google.adk.flows.llmflows.ResumabilityConfig;
import com.google.adk.memory.BaseMemoryService;
import com.google.adk.models.LlmCallsLimitExceededException;
import com.google.adk.plugins.Plugin;
import com.google.adk.plugins.PluginManager;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
Expand All @@ -42,7 +43,7 @@ public class InvocationContext {
private final BaseSessionService sessionService;
private final BaseArtifactService artifactService;
private final BaseMemoryService memoryService;
private final PluginManager pluginManager;
private final Plugin pluginManager;
private final Optional<LiveRequestQueue> liveRequestQueue;
private final Map<String, ActiveStreamingTool> activeStreamingTools = new ConcurrentHashMap<>();
private final String invocationId;
Expand Down Expand Up @@ -80,7 +81,7 @@ public InvocationContext(
BaseSessionService sessionService,
BaseArtifactService artifactService,
BaseMemoryService memoryService,
PluginManager pluginManager,
Plugin pluginManager,
Optional<LiveRequestQueue> liveRequestQueue,
Optional<String> branch,
String invocationId,
Expand Down Expand Up @@ -235,7 +236,7 @@ public BaseMemoryService memoryService() {
}

/** Returns the plugin manager for accessing tools and plugins. */
public PluginManager pluginManager() {
public Plugin pluginManager() {
return pluginManager;
}

Expand Down Expand Up @@ -376,7 +377,7 @@ public static class Builder {
private BaseSessionService sessionService;
private BaseArtifactService artifactService;
private BaseMemoryService memoryService;
private PluginManager pluginManager = new PluginManager();
private Plugin pluginManager = new PluginManager();
private Optional<LiveRequestQueue> liveRequestQueue = Optional.empty();
private Optional<String> branch = Optional.empty();
private String invocationId = newInvocationContextId();
Expand Down Expand Up @@ -430,7 +431,7 @@ public Builder memoryService(BaseMemoryService memoryService) {
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
public Builder pluginManager(PluginManager pluginManager) {
public Builder pluginManager(Plugin pluginManager) {
this.pluginManager = pluginManager;
return this;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ private Flowable<LlmResponse> callLlm(
exception ->
context
.pluginManager()
.runOnModelErrorCallback(
.onModelErrorCallback(
new CallbackContext(
context, eventForCallbackUsage.actions()),
llmRequestBuilder,
Expand Down Expand Up @@ -244,7 +244,7 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());

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

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

Expand Down Expand Up @@ -280,7 +280,7 @@ private Single<LlmResponse> handleAfterModelCallback(
CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());

Maybe<LlmResponse> pluginResult =
context.pluginManager().runAfterModelCallback(callbackContext, llmResponse);
context.pluginManager().afterModelCallback(callbackContext, llmResponse);

LlmAgent agent = (LlmAgent) context.agent();
Optional<List<? extends AfterModelCallback>> callbacksOpt = agent.afterModelCallback();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,7 @@ private static Maybe<Event> postProcessFunctionResult(
t ->
invocationContext
.pluginManager()
.runOnToolErrorCallback(tool, functionArgs, toolContext, t)
.onToolErrorCallback(tool, functionArgs, toolContext, t)
.map(isLive ? Optional::ofNullable : Optional::of)
.switchIfEmpty(Single.error(t)))
.flatMapMaybe(
Expand Down Expand Up @@ -462,7 +462,7 @@ private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
LlmAgent agent = (LlmAgent) invocationContext.agent();

Maybe<Map<String, Object>> pluginResult =
invocationContext.pluginManager().runBeforeToolCallback(tool, functionArgs, toolContext);
invocationContext.pluginManager().beforeToolCallback(tool, functionArgs, toolContext);

Optional<List<? extends BeforeToolCallback>> callbacksOpt = agent.beforeToolCallback();
if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
Expand Down Expand Up @@ -496,7 +496,7 @@ private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
Maybe<Map<String, Object>> pluginResult =
invocationContext
.pluginManager()
.runAfterToolCallback(tool, functionArgs, toolContext, functionResult);
.afterToolCallback(tool, functionArgs, toolContext, functionResult);

Optional<List<? extends AfterToolCallback>> callbacksOpt = agent.afterToolCallback();
if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
Expand Down
170 changes: 2 additions & 168 deletions core/src/main/java/com/google/adk/plugins/BasePlugin.java
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,6 @@
*/
package com.google.adk.plugins;

import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.InvocationContext;
import com.google.adk.events.Event;
import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.genai.types.Content;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import java.util.Map;

/**
* Base class for creating plugins.
*
Expand All @@ -40,168 +27,15 @@
* <p>A plugin can implement one or more methods of callbacks, but should not implement the same
* method of callback for multiple times.
*/
public abstract class BasePlugin {
public abstract class BasePlugin implements Plugin {
protected final String name;

public BasePlugin(String name) {
this.name = name;
}

@Override
public String getName() {
return name;
}

/**
* Callback executed when a user message is received before an invocation starts.
*
* @param invocationContext The context for the entire invocation.
* @param userMessage The message content input by user.
* @return An optional Content to replace the user message. Returning Empty to proceed normally.
*/
public Maybe<Content> onUserMessageCallback(
InvocationContext invocationContext, Content userMessage) {
return Maybe.empty();
}

/**
* Callback executed before the ADK runner runs.
*
* @param invocationContext The context for the entire invocation.
* @return An optional Content to halt execution. Returning Empty to proceed normally.
*/
public Maybe<Content> beforeRunCallback(InvocationContext invocationContext) {
return Maybe.empty();
}

/**
* Callback executed after an event is yielded from runner.
*
* @param invocationContext The context for the entire invocation.
* @param event The event raised by the runner.
* @return An optional Event to modify or replace the response. Returning Empty to proceed
* normally.
*/
public Maybe<Event> onEventCallback(InvocationContext invocationContext, Event event) {
return Maybe.empty();
}

/**
* Callback executed after an ADK runner run has completed.
*
* @param invocationContext The context for the entire invocation.
*/
public Completable afterRunCallback(InvocationContext invocationContext) {
return Completable.complete();
}

/**
* Callback executed before an agent's primary logic is invoked.
*
* @param agent The agent that is about to run.
* @param callbackContext The context for the agent invocation.
* @return An optional Content object to bypass the agent's execution. Returning Empty to proceed
* normally.
*/
public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
return Maybe.empty();
}

/**
* Callback executed after an agent's primary logic has completed.
*
* @param agent The agent that has just run.
* @param callbackContext The context for the agent invocation.
* @return An optional Content object to replace the agent's original result. Returning Empty to
* use the original result.
*/
public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
return Maybe.empty();
}

/**
* Callback executed before a request is sent to the model.
*
* @param callbackContext The context for the current agent call.
* @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.Builder llmRequest) {
return Maybe.empty();
}

/**
* Callback executed after a response is received from the model.
*
* @param callbackContext The context for the current agent call.
* @param llmResponse The response object received from the model.
* @return An optional LlmResponse to modify or replace the response. Returning Empty to use the
* original response.
*/
public Maybe<LlmResponse> afterModelCallback(
CallbackContext callbackContext, LlmResponse llmResponse) {
return Maybe.empty();
}

/**
* Callback executed when a model call encounters an error.
*
* @param callbackContext The context for the current agent call.
* @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.Builder llmRequest, Throwable error) {
return Maybe.empty();
}

/**
* Callback executed before a tool is called.
*
* @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.
* @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) {
return Maybe.empty();
}

/**
* Callback executed after a tool has been called.
*
* @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 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,
Map<String, Object> result) {
return Maybe.empty();
}

/**
* Callback executed when a tool call encounters an error.
*
* @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 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) {
return Maybe.empty();
}
}
Loading