4848import com .google .adk .models .Model ;
4949import com .google .adk .tools .BaseTool ;
5050import com .google .adk .tools .BaseToolset ;
51+ import com .google .adk .tools .ToolMarker ;
5152import com .google .common .base .Preconditions ;
5253import com .google .common .collect .ImmutableList ;
5354import com .google .errorprone .annotations .CanIgnoreReturnValue ;
@@ -84,7 +85,7 @@ public enum IncludeContents {
8485 private final Optional <Model > model ;
8586 private final Instruction instruction ;
8687 private final Instruction globalInstruction ;
87- private final List <Object > toolsUnion ;
88+ private final List <ToolMarker > toolsUnion ;
8889 private final ImmutableList <BaseToolset > toolsets ;
8990 private final Optional <GenerateContentConfig > generateContentConfig ;
9091 // TODO: Remove exampleProvider field - examples should only be provided via ExampleTool
@@ -152,7 +153,7 @@ public static Builder builder() {
152153 }
153154
154155 /** Extracts BaseToolset instances from the toolsUnion list. */
155- private static ImmutableList <BaseToolset > extractToolsets (List <Object > toolsUnion ) {
156+ private static ImmutableList <BaseToolset > extractToolsets (List <ToolMarker > toolsUnion ) {
156157 return toolsUnion .stream ()
157158 .filter (obj -> obj instanceof BaseToolset )
158159 .map (obj -> (BaseToolset ) obj )
@@ -165,7 +166,7 @@ public static class Builder extends BaseAgent.Builder<Builder> {
165166
166167 private Instruction instruction ;
167168 private Instruction globalInstruction ;
168- private ImmutableList <Object > toolsUnion ;
169+ private ImmutableList <ToolMarker > toolsUnion ;
169170 private GenerateContentConfig generateContentConfig ;
170171 private BaseExampleProvider exampleProvider ;
171172 private IncludeContents includeContents ;
@@ -221,13 +222,13 @@ public Builder globalInstruction(String globalInstruction) {
221222 }
222223
223224 @ CanIgnoreReturnValue
224- public Builder tools (List <?> tools ) {
225+ public Builder tools (List <? extends ToolMarker > tools ) {
225226 this .toolsUnion = ImmutableList .copyOf (tools );
226227 return this ;
227228 }
228229
229230 @ CanIgnoreReturnValue
230- public Builder tools (Object ... tools ) {
231+ public Builder tools (ToolMarker ... tools ) {
231232 this .toolsUnion = ImmutableList .copyOf (tools );
232233 return this ;
233234 }
@@ -679,7 +680,7 @@ public Single<Map.Entry<String, Boolean>> canonicalGlobalInstruction(ReadonlyCon
679680 */
680681 public Flowable <BaseTool > canonicalTools (Optional <ReadonlyContext > context ) {
681682 List <Flowable <BaseTool >> toolFlowables = new ArrayList <>();
682- for (Object toolOrToolset : toolsUnion ) {
683+ for (ToolMarker toolOrToolset : toolsUnion ) {
683684 if (toolOrToolset instanceof BaseTool baseTool ) {
684685 toolFlowables .add (Flowable .just (baseTool ));
685686 } else if (toolOrToolset instanceof BaseToolset baseToolset ) {
@@ -740,7 +741,7 @@ public Single<List<BaseTool>> tools() {
740741 return canonicalTools ().toList ();
741742 }
742743
743- public List <Object > toolsUnion () {
744+ public List <ToolMarker > toolsUnion () {
744745 return toolsUnion ;
745746 }
746747
0 commit comments