Fix mypy errors

Author: holtskinner

examples/google_adk/calendar_agent/adk_agent_executor.py

@@ -8,6 +8,9 @@
 from google.adk import Runner
 from google.adk.auth import AuthConfig
 from google.adk.events import Event
+from google.adk.session_service import (
+    Session,
+)
 from google.genai import types
 
 from a2a.server.agent_execution import AgentExecutor
@@ -44,6 +47,7 @@ class ADKAgentExecutor(AgentExecutor):
     """An AgentExecutor that runs an ADK-based Agent."""
 
     _awaiting_auth: dict[str, asyncio.Future]
+    _running_sessions: dict[str, Session]
 
     def __init__(self, runner: Runner, card: AgentCard):
         self.runner = runner
@@ -53,7 +57,7 @@ def __init__(self, runner: Runner, card: AgentCard):
 
     def _run_agent(
         self, session_id, new_message: types.Content
-    ) -> AsyncGenerator[Event, None]:
+    ) -> AsyncGenerator[Event]:
         return self.runner.run_async(
             session_id=session_id, user_id='self', new_message=new_message
         )
@@ -97,12 +101,17 @@ async def _process_request(
                 # is received.
                 break
             if event.is_final_response():
-                parts = convert_genai_parts_to_a2a(event.content.parts)
-                logger.debug('Yielding final response: %s', parts)
-                task_updater.add_artifact(parts)
-                task_updater.complete()
+                if event.content and event.content.parts:
+                    parts = convert_genai_parts_to_a2a(event.content.parts)
+                    logger.debug('Yielding final response: %s', parts)
+                    task_updater.add_artifact(parts)
+                    task_updater.complete()
                 break
-            if not event.get_function_calls():
+            if (
+                not event.get_function_calls()
+                and event.content
+                and event.content.parts
+            ):
                 logger.debug('Yielding update response')
                 task_updater.update_status(
                     TaskState.working,
@@ -228,12 +237,17 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue):
     async def on_auth_callback(self, state: str, uri: str):
         self._awaiting_auth[state].set_result(uri)
 
-    def _upsert_session(self, session_id: str):
-        return self.runner.session_service.get_session(
-            app_name=self.runner.app_name, user_id='self', session_id=session_id
-        ) or self.runner.session_service.create_session(
+    def _upsert_session(self, session_id: str) -> Session:
+        session = self.runner.session_service.get_session(
             app_name=self.runner.app_name, user_id='self', session_id=session_id
         )
+        if session is None:
+            session = self.runner.session_service.create_session(
+                app_name=self.runner.app_name,
+                user_id='self',
+                session_id=session_id,
+            )
+        return session
 
 
 def convert_a2a_parts_to_genai(parts: list[Part]) -> list[types.Part]:
@@ -243,32 +257,43 @@ def convert_a2a_parts_to_genai(parts: list[Part]) -> list[types.Part]:
 
 def convert_a2a_part_to_genai(part: Part) -> types.Part:
     """Convert a single A2A Part type into a Google Gen AI Part type."""
-    part = part.root
-    if isinstance(part, TextPart):
-        return types.Part(text=part.text)
-    if isinstance(part, FilePart):
-        if isinstance(part.file, FileWithUri):
+    if isinstance(part.root, TextPart):
+        return types.Part(text=part.root.text)
+    if isinstance(part.root, FilePart):
+        file_data = part.root.file
+        if isinstance(file_data, FileWithUri):
             return types.Part(
                 file_data=types.FileData(
-                    file_uri=part.file.uri, mime_type=part.file.mime_type
+                    file_uri=file_data.uri,
+                    mime_type=file_data.mime_type,
                 )
             )
-        if isinstance(part.file, FileWithBytes):
+        if isinstance(file_data, FileWithBytes):
             return types.Part(
                 inline_data=types.Blob(
-                    data=part.file.bytes, mime_type=part.file.mime_type
+                    data=file_data.bytes,
+                    mime_type=file_data.mime_type,
                 )
             )
-        raise ValueError(f'Unsupported file type: {type(part.file)}')
-    raise ValueError(f'Unsupported part type: {type(part)}')
+        raise ValueError(f'Unsupported file type: {type(file_data)}')
+    raise ValueError(f'Unsupported part root type: {type(part.root)}')
 
 
 def convert_genai_parts_to_a2a(parts: list[types.Part]) -> list[Part]:
     """Convert a list of Google Gen AI Part types into a list of A2A Part types."""
+    if not parts:
+        return []
     return [
         convert_genai_part_to_a2a(part)
         for part in parts
-        if (part.text or part.file_data or part.inline_data)
+        if part
+        and (
+            part.text
+            or part.file_data
+            or part.inline_data
+            or part.function_call
+            or part.function_response
+        )
     ]
 
 
@@ -295,7 +320,9 @@ def convert_genai_part_to_a2a(part: types.Part) -> Part:
     raise ValueError(f'Unsupported part type: {part}')
 
 
-def get_auth_request_function_call(event: Event) -> types.FunctionCall:
+def get_auth_request_function_call(
+    event: Event,
+) -> types.FunctionCall | None:
     """Get the special auth request function call from the event."""
     if not (event.content and event.content.parts):
         return None
@@ -315,8 +342,12 @@ def get_auth_config(
     auth_request_function_call: types.FunctionCall,
 ) -> AuthConfig:
     """Extracts the AuthConfig object from the arguments of the auth request function call."""
-    if not auth_request_function_call.args or not (
-        auth_config := auth_request_function_call.args.get('auth_config')
+    if (
+        not auth_request_function_call.args
+        or not isinstance(auth_request_function_call.args, dict)
+        or not (
+            auth_config := auth_request_function_call.args.get('auth_config')
+        )
     ):
         raise ValueError(
             f'Cannot get auth config from function call: {auth_request_function_call}'

src/a2a/client/client.py

@@ -133,7 +133,7 @@ async def send_message_streaming(
         request: SendStreamingMessageRequest,
         *,
         http_kwargs: dict[str, Any] | None = None,
-    ) -> AsyncGenerator[SendStreamingMessageResponse, None]:
+    ) -> AsyncGenerator[SendStreamingMessageResponse]:
         if not request.id:
             request.id = str(uuid4())