chore: initial fix of mypy errors in vertexai/_genai

PiperOrigin-RevId: 778192717

Author: sararob

.github/workflows/mypy.yml

@@ -0,0 +1,37 @@
+# This workflow runs mypy for static type checking on the vertexai/_genai submodule.
+# See https://mypy.readthedocs.io/en/stable/index.html for more information.
+#
+# You can adjust the behavior by modifying this file.
+name: Run mypy
+
+on:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - 'vertexai/_genai/**'
+
+jobs:
+  genai-mypy:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11', '3.12']
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install mypy
+          pip install google-cloud-aiplatform
+
+      - name: Run mypy  ${{ matrix.python-version }}
+        run: mypy vertexai/_genai/ --strict --config-file=mypy.ini

mypy.ini

@@ -1,3 +1,22 @@
 [mypy]
-python_version = 3.7
-namespace_packages = True
+# TODO(b/422425982): Fix arg-type errors
+disable_error_code = import-not-found, import-untyped, arg-type
+
+# We only want to run mypy on _genai dir, ignore dependent modules
+[mypy-vertexai.agent_engines.*]
+ignore_errors = True
+
+[mypy-vertexai.preview.*]
+ignore_errors = True
+
+[mypy-vertexai.generative_models.*]
+ignore_errors = True
+
+[mypy-vertexai.prompts.*]
+ignore_errors = True
+
+[mypy-vertexai.tuning.*]
+ignore_errors = True
+
+[mypy-vertexai.caching.*]
+ignore_errors = True

vertexai/__init__.py

@@ -27,7 +27,7 @@
 _genai_types = None
 
 
-def __getattr__(name):
+def __getattr__(name):  # type: ignore[no-untyped-def]
     # Lazy importing the preview submodule
     # See https://peps.python.org/pep-0562/
     if name == "preview":

vertexai/_genai/__init__.py

@@ -16,12 +16,12 @@
 
 import importlib
 
-from .client import Client
+from .client import Client  # type: ignore[attr-defined]
 
 _evals = None
 
 
-def __getattr__(name):
+def __getattr__(name):  # type: ignore[no-untyped-def]
     if name == "evals":
         global _evals
         if _evals is None:

vertexai/_genai/_agent_engines_utils.py

@@ -34,7 +34,7 @@ def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
         the `query` API.
     """
 
-    def _method(self: types.AgentEngine, **kwargs):
+    def _method(self: types.AgentEngine, **kwargs) -> Any:
         if not self.api_client:
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
@@ -52,7 +52,9 @@ def _method(self: types.AgentEngine, **kwargs):
     return _method
 
 
-def _wrap_async_query_operation(*, method_name: str) -> Callable[..., Coroutine]:
+def _wrap_async_query_operation(
+    *, method_name: str
+) -> Callable[..., Coroutine[Any, Any, Any]]:
     """Wraps an Agent Engine method, creating an async callable for `query` API.
 
     This function creates a callable object that executes the specified
@@ -122,7 +124,7 @@ def _method(self: types.AgentEngine, **kwargs) -> Iterator[Any]:
 
 def _wrap_async_stream_query_operation(
     *, method_name: str
-) -> Callable[..., AsyncIterator]:
+) -> Callable[..., AsyncIterator[Any]]:
     """Wraps an Agent Engine method, creating an async callable for `stream_query` API.
 
     This function creates a callable object that executes the specified

vertexai/_genai/_evals_common.py

@@ -253,7 +253,7 @@ def _execute_inference_concurrently(
                         e,
                     )
                     responses[index] = {"error": f"Inference task failed: {e}"}
-    return responses
+    return responses  # type: ignore[return-value]
 
 
 def _run_gemini_inference(
@@ -337,14 +337,14 @@ def _call_litellm_completion(model: str, messages: list[dict[str, Any]]) -> dict
 
 def _run_litellm_inference(
     model: str, prompt_dataset: pd.DataFrame
-) -> list[dict[str, Any]]:
+) -> list[Optional[dict[str, Any]]]:
     """Runs inference using LiteLLM with concurrency."""
     logger.info(
         "Generating responses for %d prompts using LiteLLM for third party model: %s",
         len(prompt_dataset),
         model,
     )
-    responses = [None] * len(prompt_dataset)
+    responses: list[Optional[dict[str, Any]]] = [None] * len(prompt_dataset)
     tasks = []
 
     with tqdm(total=len(prompt_dataset), desc=f"LiteLLM Inference ({model})") as pbar:
@@ -714,7 +714,7 @@ def _resolve_dataset_inputs(
         loaded_raw_datasets.append(current_loaded_data)
 
         if dataset_schema:
-            current_schema = _evals_data_converters.EvalDatasetSchema[dataset_schema]
+            current_schema = _evals_data_converters.EvalDatasetSchema(dataset_schema)
         else:
             current_schema = _evals_data_converters.auto_detect_dataset_schema(
                 current_loaded_data
@@ -826,11 +826,13 @@ def _execute_evaluation(
             " EvaluationDataset or a list of EvaluationDataset."
         )
     original_candidate_names = [
-        ds.candidate_name or f"candidate_{i+1}" for i, ds in enumerate(dataset_list)
+        ds.candidate_name or f"candidate_{i + 1}" for i, ds in enumerate(dataset_list)
     ]
     name_counts = collections.Counter(original_candidate_names)
     deduped_candidate_names = []
-    current_name_counts = collections.defaultdict(int)
+    current_name_counts: collections.defaultdict[Any, int] = collections.defaultdict(
+        int
+    )
 
     for name in original_candidate_names:
         if name_counts[name] > 1:

vertexai/_genai/_evals_data_converters.py

@@ -17,7 +17,7 @@
 import abc
 import json
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from google.genai import _common
 from google.genai import types as genai_types
@@ -103,12 +103,20 @@ def _parse_request(
                 last_message.content.role if last_message.content else "user"
             )
             if last_message_role in ["user", None]:
-                prompt = last_message.content
+                prompt = (
+                    last_message.content
+                    if last_message.content
+                    else genai_types.Content()
+                )
             elif last_message_role == "model":
                 reference = types.ResponseCandidate(response=last_message.content)
                 if conversation_history:
                     second_to_last_message = conversation_history.pop()
-                    prompt = second_to_last_message.content
+                    prompt = (
+                        second_to_last_message.content
+                        if second_to_last_message.content
+                        else genai_types.Content()
+                    )
                 else:
                     prompt = genai_types.Content()
 
@@ -436,7 +444,7 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset:
 
 def auto_detect_dataset_schema(
     raw_dataset: list[dict[str, Any]],
-) -> EvalDatasetSchema:
+) -> Union[EvalDatasetSchema, str]:
     """Detects the schema of a raw dataset."""
     if not raw_dataset:
         return EvalDatasetSchema.UNKNOWN
@@ -522,7 +530,7 @@ def _validate_case_consistency(
     current_case: types.EvalCase,
     case_idx: int,
     dataset_idx: int,
-):
+) -> None:
     """Logs warnings if prompt or reference mismatches occur."""
     if base_case.prompt != current_case.prompt:
         base_prompt_text_preview = _get_first_part_text(base_case.prompt)[:50]
@@ -609,7 +617,11 @@ def merge_response_datasets_into_canonical_format(
     base_parsed_dataset = parsed_evaluation_datasets[0]
 
     for case_idx in range(num_expected_cases):
-        base_eval_case: types.EvalCase = base_parsed_dataset.eval_cases[case_idx]
+        base_eval_case: types.EvalCase = (
+            base_parsed_dataset.eval_cases[case_idx]
+            if base_parsed_dataset.eval_cases
+            else types.EvalCase()
+        )
         candidate_responses: list[types.ResponseCandidate] = []
 
         if base_eval_case.responses:
@@ -640,9 +652,11 @@ def merge_response_datasets_into_canonical_format(
         for dataset_idx_offset, current_parsed_ds in enumerate(
             parsed_evaluation_datasets[1:], start=1
         ):
-            current_ds_eval_case: types.EvalCase = current_parsed_ds.eval_cases[
-                case_idx
-            ]
+            current_ds_eval_case: types.EvalCase = (
+                current_parsed_ds.eval_cases[case_idx]
+                if current_parsed_ds.eval_cases
+                else types.EvalCase()
+            )
 
             _validate_case_consistency(
                 base_eval_case, current_ds_eval_case, case_idx, dataset_idx_offset