chore: Gen AI SDK client - resolve some agent engine mypy errors

PiperOrigin-RevId: 833396344

Author: sararob

vertexai/_genai/_agent_engines_utils.py

@@ -64,23 +64,23 @@
     _BUILTIN_MODULE_NAMES: Sequence[str] = []  # type: ignore[no-redef]
 
 try:
-    _PACKAGE_DISTRIBUTIONS: Mapping[
-        str, Sequence[str]
-    ] = importlib_metadata.packages_distributions()  # type: ignore[attr-defined]
+    _PACKAGE_DISTRIBUTIONS: Mapping[str, Sequence[str]] = (
+        importlib_metadata.packages_distributions()
+    )
 except AttributeError:
     _PACKAGE_DISTRIBUTIONS: Mapping[str, Sequence[str]] = {}  # type: ignore[no-redef]
 
 try:
     # sys.stdlib_module_names is available from Python 3.10 onwards.
-    _STDLIB_MODULE_NAMES: frozenset[str] = sys.stdlib_module_names  # type: ignore[attr-defined]
+    _STDLIB_MODULE_NAMES: frozenset[str] = sys.stdlib_module_names
 except AttributeError:
     _STDLIB_MODULE_NAMES: frozenset[str] = frozenset()  # type: ignore[no-redef]
 
 
 try:
-    from google.cloud import storage  # type: ignore[attr-defined]
+    from google.cloud import storage
 
-    _StorageBucket: type[Any] = storage.Bucket
+    _StorageBucket: type[Any] = storage.Bucket  # type: ignore[attr-defined]
 except (ImportError, AttributeError):
     _StorageBucket: type[Any] = Any  # type: ignore[no-redef]
 
@@ -90,7 +90,7 @@
 
     _SpecifierSet: type[Any] = packaging.specifiers.SpecifierSet
 except (ImportError, AttributeError):
-    _SpecifierSet: type[Any] = Any
+    _SpecifierSet: type[Any] = Any  # type: ignore[no-redef]
 
 
 try:
@@ -125,7 +125,7 @@
 _ASYNC_API_MODE = "async"
 _ASYNC_STREAM_API_MODE = "async_stream"
 _BIDI_STREAM_API_MODE = "bidi_stream"
-_BASE_MODULES = set(_BUILTIN_MODULE_NAMES + tuple(_STDLIB_MODULE_NAMES))
+_BASE_MODULES = set(_BUILTIN_MODULE_NAMES).union(_STDLIB_MODULE_NAMES)
 _BLOB_FILENAME = "agent_engine.pkl"
 _DEFAULT_AGENT_FRAMEWORK = "custom"
 _SUPPORTED_AGENT_FRAMEWORKS = frozenset(
@@ -237,7 +237,9 @@ class BidiStreamQueryable(Protocol):
     """Protocol for Agent Engines that can stream requests and responses."""
 
     @abc.abstractmethod
-    async def bidi_stream_query(self, input_queue: asyncio.Queue) -> AsyncIterator[Any]:
+    async def bidi_stream_query(
+        self, input_queue: asyncio.Queue[Any]
+    ) -> AsyncIterator[Any]:
         """Stream requests and responses to serve the user queries."""
 
 
@@ -255,7 +257,7 @@ class OperationRegistrable(Protocol):
     """Protocol for agents that have registered operations."""
 
     @abc.abstractmethod
-    def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:
+    def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:  # type: ignore[no-untyped-def]
         """Register the user provided operations (modes and methods)."""
 
 
@@ -264,7 +266,7 @@ def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:
 
     ADKAgent: type[Any] = BaseAgent
 except (ImportError, AttributeError):
-    ADKAgent: type[Any] = Any
+    ADKAgent: type[Any] = Any  # type: ignore[no-redef]
 
 _AgentEngineInterface = Union[
     ADKAgent,
@@ -399,13 +401,15 @@ class _RequirementsValidationResult(TypedDict):
 
 
 class GetOperationFunction(Protocol):
-    def __call__(self, *, operation_name: str, **kwargs) -> AgentEngineOperationUnion:
+    def __call__(
+        self, *, operation_name: str, **kwargs: Any
+    ) -> AgentEngineOperationUnion:
         pass
 
 
 class GetAsyncOperationFunction(Protocol):
     async def __call__(
-        self, *, operation_name: str, **kwargs
+        self, *, operation_name: str, **kwargs: Any
     ) -> Awaitable[AgentEngineOperationUnion]:
         pass
 
@@ -727,7 +731,7 @@ def _get_agent_framework(
     *,
     agent_framework: Optional[str],
     agent: _AgentEngineInterface,
-) -> str:
+) -> Union[str, Any]:
     """Gets the agent framework to use.
 
     The agent framework is determined in the following order of priority:
@@ -820,13 +824,13 @@ def _import_cloudpickle_or_raise() -> types.ModuleType:
 def _import_cloud_storage_or_raise() -> types.ModuleType:
     """Tries to import the Cloud Storage module."""
     try:
-        from google.cloud import storage  # type: ignore[attr-defined]
+        from google.cloud import storage
     except ImportError as e:
         raise ImportError(
             "Cloud Storage is not installed. Please call "
             "'pip install google-cloud-aiplatform[agent_engines]'."
         ) from e
-    return storage  # type: ignore[no-any-return]
+    return storage
 
 
 def _import_packaging_requirements_or_raise() -> types.ModuleType:
@@ -838,7 +842,7 @@ def _import_packaging_requirements_or_raise() -> types.ModuleType:
             "packaging.requirements is not installed. Please call "
             "'pip install google-cloud-aiplatform[agent_engines]'."
         ) from e
-    return requirements  # type: ignore[no-any-return]
+    return requirements
 
 
 def _import_packaging_version_or_raise() -> types.ModuleType:
@@ -850,7 +854,7 @@ def _import_packaging_version_or_raise() -> types.ModuleType:
             "packaging.version is not installed. Please call "
             "'pip install google-cloud-aiplatform[agent_engines]'."
         ) from e
-    return version  # type: ignore[no-any-return]
+    return version
 
 
 def _import_pydantic_or_raise() -> types.ModuleType:
@@ -866,7 +870,7 @@ def _import_pydantic_or_raise() -> types.ModuleType:
             "pydantic is not installed. Please call "
             "'pip install google-cloud-aiplatform[agent_engines]'."
         ) from e
-    return pydantic  # type: ignore[no-any-return]
+    return pydantic
 
 
 def _parse_constraints(
@@ -1022,7 +1026,7 @@ def _register_api_methods_or_raise(
         }
         if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn:
             # Override the default function with user-specified function if it exists.
-            _wrap_operation = wrap_operation_fn[api_mode]  # type: ignore[assignment]
+            _wrap_operation = wrap_operation_fn[api_mode]
         elif api_mode in _wrap_operation_map:
             _wrap_operation = _wrap_operation_map[api_mode]  # type: ignore[assignment]
         else:
@@ -1531,7 +1535,7 @@ def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
 
 class GetOperationFunction(Protocol):
     def __call__(  # noqa: E704
-        self, *, operation_name: str, **kwargs
+        self, *, operation_name: str, **kwargs: Any
     ) -> AgentEngineOperationUnion: ...
 
 
@@ -1588,8 +1592,8 @@ def _wrap_async_query_operation(
     """
 
     async def _method(
-        self: genai_types.AgentEngine, **kwargs
-    ) -> Coroutine[Any, Any, Any]:  # type: ignore[no-untyped-def]
+        self: genai_types.AgentEngine, **kwargs: Any
+    ) -> Union[Coroutine[Any, Any, Any], Any]:
         if not self.api_async_client:
             raise ValueError("api_async_client is not initialized.")
         if not self.api_resource:
@@ -1681,7 +1685,7 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
     return _method
 
 
-def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]:
+def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
     """Wraps an Agent Engine method, creating a callable for A2A API.
 
     Args:
@@ -1727,7 +1731,7 @@ def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list
         the A2A API.
     """
 
-    async def _method(self, **kwargs) -> Any:
+    async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
         """Wraps an Agent Engine method, creating a callable for A2A API."""
         if not self.api_client:
             raise ValueError("api_client is not initialized.")

vertexai/_genai/agent_engines.py

@@ -703,7 +703,7 @@ def _update(
     _sessions = None
 
     @property
-    def memories(self):
+    def memories(self) -> Any:
         if self._memories is None:
             try:
                 # We need to lazy load the memories module to handle the
@@ -722,7 +722,7 @@ def memories(self):
         "The Vertex SDK GenAI agent_engines.sandboxes module is experimental, "
         "and may change in future versions."
     )
-    def sandboxes(self):
+    def sandboxes(self) -> Any:
         if self._sandboxes is None:
             try:
                 # We need to lazy load the sandboxes module to handle the
@@ -737,7 +737,7 @@ def sandboxes(self):
         return self._sandboxes.Sandboxes(self._api_client)
 
     @property
-    def sessions(self):
+    def sessions(self) -> Any:
         if self._sessions is None:
             try:
                 # We need to lazy load the sessions module to handle the
@@ -1039,7 +1039,7 @@ def _create_config(
             sys_version = python_version
         else:
             sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-        agent_engine_spec = None
+        agent_engine_spec: Any = None
         if agent is not None:
             if source_packages is not None:
                 raise ValueError(