chore: fix mypy strict type errors in adk agents

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 945259368

Author: GWeale

src/google/adk/agents/active_streaming_tool.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import asyncio
+from typing import Any
 from typing import Optional
 
 from pydantic import BaseModel
@@ -32,7 +33,7 @@ class ActiveStreamingTool(BaseModel):
   )
   """The pydantic model config."""
 
-  task: Optional[asyncio.Task] = None
+  task: Optional[asyncio.Task[Any]] = None
   """The active task of this streaming tool."""
 
   stream: Optional[LiveRequestQueue] = None

src/google/adk/agents/base_agent.py

@@ -581,7 +581,7 @@ def model_post_init(self, __context: Any) -> None:
 
   @field_validator('name', mode='after')
   @classmethod
-  def validate_name(cls, value: str):
+  def validate_name(cls, value: str) -> str:
     if not value.isidentifier():
       raise ValueError(
           f'Found invalid agent name: `{value}`.'

src/google/adk/agents/config_agent_utils.py

@@ -264,7 +264,7 @@ def resolve_agent_reference(
     raise ValueError("AgentRefConfig must have either 'code' or 'config_path'")
 
 
-def _resolve_agent_code_reference(code: str) -> Any:
+def _resolve_agent_code_reference(code: str) -> BaseAgent:
   """Resolve a code reference to an actual agent instance.
 
   Args:

src/google/adk/agents/context.py

@@ -86,8 +86,8 @@ def _derive_node_path(
 
   if not parent_path and isinstance(node, BaseAgent) and node.parent_agent:
     path_builder = _NodePathBuilder([])
-    curr = node.parent_agent
-    parent_agents = []
+    curr: BaseAgent | None = node.parent_agent
+    parent_agents: list[BaseAgent] = []
     depth = 0
     while curr is not None and depth < _MAX_PARENT_DEPTH:
       parent_agents.insert(0, curr)

src/google/adk/agents/invocation_context.py

@@ -16,7 +16,6 @@
 
 import asyncio
 from typing import Any
-from typing import cast
 from typing import Optional
 
 from google.adk.platform import uuid as platform_uuid
@@ -84,7 +83,7 @@ class _InvocationCostManager(BaseModel):
 
   def increment_and_enforce_llm_calls_limit(
       self, run_config: Optional[RunConfig]
-  ):
+  ) -> None:
     """Increments _number_of_llm_calls and enforces the limit."""
     # We first increment the counter and then check the conditions.
     self._number_of_llm_calls += 1
@@ -389,7 +388,7 @@ def populate_invocation_agent_states(self) -> None:
 
   def increment_llm_call_count(
       self,
-  ):
+  ) -> None:
     """Tracks number of llm calls made.
 
     Raises:
@@ -525,4 +524,4 @@ def stamp_event_branch_context(self, event: Event) -> None:
 
 
 def new_invocation_context_id() -> str:
-  return "e-" + cast(str, platform_uuid.new_uuid())
+  return "e-" + platform_uuid.new_uuid()

src/google/adk/agents/langgraph_agent.py

@@ -19,6 +19,7 @@
 
 from google.genai import types
 from langchain_core.messages import AIMessage
+from langchain_core.messages import BaseMessage
 from langchain_core.messages import HumanMessage
 from langchain_core.messages import SystemMessage
 from langchain_core.runnables.config import RunnableConfig
@@ -31,7 +32,9 @@
 from .invocation_context import InvocationContext
 
 
-def _get_last_human_messages(events: list[Event]) -> list[HumanMessage]:
+def _get_last_human_messages(
+    events: list[Event],
+) -> list[Union[HumanMessage, AIMessage]]:
   """Extracts last human messages from given list of events.
 
   Args:
@@ -40,7 +43,7 @@ def _get_last_human_messages(events: list[Event]) -> list[HumanMessage]:
   Returns:
     list of last human messages
   """
-  messages = []
+  messages: list[Union[HumanMessage, AIMessage]] = []
   for event in reversed(events):
     if messages and event.author != 'user':
       break
@@ -77,7 +80,7 @@ async def _run_async_impl(
         if current_graph_state.values
         else []
     )
-    messages = (
+    messages: list[BaseMessage] = (
         [SystemMessage(content=self.instruction)]
         if self.instruction and not graph_messages
         else []
@@ -132,7 +135,7 @@ def _get_conversation_with_agent(
       list of messages
     """
 
-    messages = []
+    messages: list[Union[HumanMessage, AIMessage]] = []
     for event in events:
       if not event.content or not event.content.parts:
         continue

src/google/adk/agents/live_request_queue.py

@@ -62,27 +62,27 @@ class LiveRequest(BaseModel):
 class LiveRequestQueue:
   """Queue used to send LiveRequest in a live(bidirectional streaming) way."""
 
-  def __init__(self):
-    self._queue = asyncio.Queue()
+  def __init__(self) -> None:
+    self._queue: asyncio.Queue[LiveRequest] = asyncio.Queue()
 
-  def close(self):
+  def close(self) -> None:
     self._queue.put_nowait(LiveRequest(close=True))
 
-  def send_content(self, content: types.Content, partial: bool = False):
+  def send_content(self, content: types.Content, partial: bool = False) -> None:
     self._queue.put_nowait(LiveRequest(content=content, partial=partial))
 
-  def send_realtime(self, blob: types.Blob):
+  def send_realtime(self, blob: types.Blob) -> None:
     self._queue.put_nowait(LiveRequest(blob=blob))
 
-  def send_activity_start(self):
+  def send_activity_start(self) -> None:
     """Sends an activity start signal to mark the beginning of user input."""
     self._queue.put_nowait(LiveRequest(activity_start=types.ActivityStart()))
 
-  def send_activity_end(self):
+  def send_activity_end(self) -> None:
     """Sends an activity end signal to mark the end of user input."""
     self._queue.put_nowait(LiveRequest(activity_end=types.ActivityEnd()))
 
-  def send(self, req: LiveRequest):
+  def send(self, req: LiveRequest) -> None:
     self._queue.put_nowait(req)
 
   async def get(self) -> LiveRequest:

src/google/adk/agents/llm_agent.py

@@ -23,7 +23,6 @@
 from typing import AsyncGenerator
 from typing import Awaitable
 from typing import Callable
-from typing import cast
 from typing import ClassVar
 from typing import Dict
 from typing import Literal
@@ -64,6 +63,7 @@
 from .base_agent import BaseAgentState
 from .base_agent_config import BaseAgentConfig as BaseAgentConfig
 from .callback_context import CallbackContext
+from .context import Context
 from .invocation_context import InvocationContext
 from .llm_agent_config import LlmAgentConfig as LlmAgentConfig
 from .readonly_context import ReadonlyContext
@@ -138,7 +138,7 @@
 
 async def _convert_tool_union_to_tools(
     tool_union: ToolUnion,
-    ctx: ReadonlyContext,
+    ctx: Optional[ReadonlyContext],
     model: Union[str, BaseLlm],
     multiple_tools: bool = False,
 ) -> list[BaseTool]:
@@ -152,7 +152,7 @@ async def _convert_tool_union_to_tools(
     from ..tools.google_search_agent_tool import create_google_search_agent
     from ..tools.google_search_agent_tool import GoogleSearchAgentTool
 
-    search_tool = cast(GoogleSearchTool, tool_union)
+    search_tool = tool_union
     if search_tool.bypass_multi_tools_limit:
       return [GoogleSearchAgentTool(create_google_search_agent(model))]
 
@@ -163,7 +163,7 @@ async def _convert_tool_union_to_tools(
   if multiple_tools and isinstance(tool_union, VertexAiSearchTool):
     from ..tools.discovery_engine_search_tool import DiscoveryEngineSearchTool
 
-    vais_tool = cast(VertexAiSearchTool, tool_union)
+    vais_tool = tool_union
     if vais_tool.bypass_multi_tools_limit:
       return [
           DiscoveryEngineSearchTool(
@@ -927,7 +927,7 @@ def _get_available_agent_names(self) -> list[str]:
     """
     agents = []
 
-    def collect_agents(agent):
+    def collect_agents(agent: BaseAgent) -> None:
       agents.append(agent.name)
       if hasattr(agent, 'sub_agents') and agent.sub_agents:
         for sub_agent in agent.sub_agents:
@@ -952,7 +952,7 @@ def __get_transfer_to_agent_or_none(
         return self.__get_agent_to_run(event.actions.transfer_to_agent)
     return None
 
-  def __maybe_save_output_to_state(self, event: Event):
+  def __maybe_save_output_to_state(self, event: Event) -> None:
     """Saves the model output to state if needed."""
     # skip if the event was authored by some other agent (e.g. current agent
     # transferred to another agent)

src/google/adk/agents/parallel_agent.py

@@ -60,7 +60,9 @@ async def _merge_agent_run(
 
   # Agents are processed in parallel.
   # Events for each agent are put on queue sequentially.
-  async def process_an_agent(events_for_one_agent):
+  async def process_an_agent(
+      events_for_one_agent: AsyncGenerator[Event, None],
+  ) -> None:
     try:
       async for event in events_for_one_agent:
         resume_signal = asyncio.Event()
@@ -112,7 +114,7 @@ async def _merge_agent_run_pre_3_11(
   sentinel = object()
   queue = asyncio.Queue()
 
-  def propagate_exceptions(tasks):
+  def propagate_exceptions(tasks: list[asyncio.Task[None]]) -> None:
     # Propagate exceptions and errors from tasks.
     for task in tasks:
       if task.done():
@@ -122,7 +124,9 @@ def propagate_exceptions(tasks):
 
   # Agents are processed in parallel.
   # Events for each agent are put on queue sequentially.
-  async def process_an_agent(events_for_one_agent):
+  async def process_an_agent(
+      events_for_one_agent: AsyncGenerator[Event, None],
+  ) -> None:
     try:
       async for event in events_for_one_agent:
         resume_signal = asyncio.Event()

src/google/adk/agents/remote_a2a_agent.py

@@ -408,7 +408,7 @@ def _create_a2a_request_for_user_function_response(
     return a2a_message
 
   def _is_remote_response(self, event: Event) -> bool:
-    return (
+    return bool(
         event.author == self.name
         and event.custom_metadata
         and event.custom_metadata.get(A2A_METADATA_PREFIX + "response", False)
@@ -448,22 +448,27 @@ def _construct_message_parts_from_session(
       events_to_process.append(event)
 
     for event in reversed(events_to_process):
+      processed_event: Optional[Event] = event
       if _is_other_agent_reply(self.name, event):
-        event = _present_other_agent_message(event)
+        processed_event = _present_other_agent_message(event)
 
-      if not event or not event.content or not event.content.parts:
+      if (
+          not processed_event
+          or not processed_event.content
+          or not processed_event.content.parts
+      ):
         continue
 
-      for part in event.content.parts:
+      for part in processed_event.content.parts:
         converted_parts = self._genai_part_converter(part)
         if not isinstance(converted_parts, list):
           converted_parts = [converted_parts] if converted_parts else []
 
-        if event.author == "user":
-          for part in converted_parts:
-            meta = _compat.part_metadata(part) or {}
+        if processed_event.author == "user":
+          for a2a_part in converted_parts:
+            meta = _compat.part_metadata(a2a_part) or {}
             meta["is_user_input"] = True
-            _compat.set_part_metadata(part, meta)
+            _compat.set_part_metadata(a2a_part, meta)
 
         if converted_parts:
           message_parts.extend(converted_parts)