chore: Resolve mypy errors in Gen AI Modules

PiperOrigin-RevId: 862235861

Author: sararob

vertexai/_genai/_agent_engines_utils.py

@@ -81,7 +81,7 @@
 try:
     from google.cloud import storage
 
-    _StorageBucket: type[Any] = storage.Bucket  # type: ignore[attr-defined]
+    _StorageBucket: type[Any] = storage.Bucket
 except (ImportError, AttributeError):
     _StorageBucket: type[Any] = Any  # type: ignore[no-redef]
 

vertexai/_genai/_evals_common.py

@@ -48,16 +48,16 @@
 try:
     import litellm
 except ImportError:
-    litellm = None  # type: ignore[assignment]
+    litellm = None
 
 try:
     from google.adk.agents import LlmAgent
     from google.adk.runners import Runner
     from google.adk.sessions import InMemorySessionService
 except ImportError:
-    LlmAgent = None  # type: ignore[assignment]
-    Runner = None  # type: ignore[assignment]
-    InMemorySessionService = None  # type: ignore[assignment]
+    LlmAgent = None
+    Runner = None
+    InMemorySessionService = None
 
 
 logger = logging.getLogger(__name__)

vertexai/_genai/_evals_metric_loaders.py

@@ -44,7 +44,7 @@ class LazyLoadedPrebuiltMetric:
         "gs://vertex-ai-generative-ai-eval-sdk-resources/metrics/{metric_name}/"
     )
 
-    def __init__(self, name: str, version: Optional[str] = None, **kwargs):
+    def __init__(self, name: str, version: Optional[str] = None, **kwargs: Any):
         self.name = name.upper()
         self.version = version
         self.metric_kwargs = kwargs
@@ -248,7 +248,7 @@ def resolve(self, api_client: Any) -> "types.Metric":
             ) from e
 
     def __call__(
-        self, version: Optional[str] = None, **kwargs
+        self, version: Optional[str] = None, **kwargs: Any
     ) -> "LazyLoadedPrebuiltMetric":
         """Allows setting a specific version and other metric attributes."""
         updated_kwargs = self.metric_kwargs.copy()
@@ -271,7 +271,7 @@ class PrebuiltMetricLoader:
     """
 
     def __getattr__(
-        self, name: str, version: Optional[str] = None, **kwargs
+        self, name: str, version: Optional[str] = None, **kwargs: Any
     ) -> LazyLoadedPrebuiltMetric:
         return LazyLoadedPrebuiltMetric(name=name, version=version, **kwargs)
 

vertexai/_genai/_evals_visualization.py

@@ -1063,7 +1063,7 @@ def display_evaluation_dataset(eval_dataset_obj: types.EvaluationDataset) -> Non
                 # Special handling for rubric_groups to keep it as a dict
                 if isinstance(cell_value, dict):
                     processed_row[col_name] = {
-                        k: [
+                        k: [  # type: ignore[misc]
                             (
                                 v_item.model_dump(mode="json")
                                 if hasattr(v_item, "model_dump")
@@ -1077,7 +1077,7 @@ def display_evaluation_dataset(eval_dataset_obj: types.EvaluationDataset) -> Non
                     processed_row[col_name] = cell_value
             else:
                 if isinstance(cell_value, (dict, list)):
-                    processed_row[col_name] = json.dumps(
+                    processed_row[col_name] = json.dumps(  # type: ignore[assignment]
                         cell_value, ensure_ascii=False, default=_pydantic_serializer
                     )
                 else:

vertexai/_genai/_gcs_utils.py

@@ -18,7 +18,7 @@
 import logging
 from typing import Any, Union
 
-from google.cloud import storage
+from google.cloud import storage  # type: ignore[attr-defined]
 from google.cloud.aiplatform.utils.gcs_utils import blob_from_uri
 from google.genai._api_client import BaseApiClient
 import pandas as pd
@@ -36,7 +36,7 @@ class GcsUtils:
 
     def __init__(self, api_client: BaseApiClient):
         self.api_client = api_client
-        self.storage_client = storage.Client(  # type: ignore[attr-defined]
+        self.storage_client = storage.Client(
             project=self.api_client.project,
             credentials=self.api_client._credentials,
         )

vertexai/_genai/_prompt_management_utils.py

@@ -99,24 +99,32 @@ def _create_prompt_from_dataset_metadata(
     api_schema = dataset.metadata.prompt_api_schema
     prompt = types.Prompt()
 
+    if api_schema is None:
+        return prompt
+
     if api_schema.multimodal_prompt:
 
         prompt_message = api_schema.multimodal_prompt.prompt_message
         prompt.prompt_data = prompt_message
 
-        executions = api_schema.executions
-        if executions:
-            prompt.prompt_data.variables = []
-            for execution in executions:
-                if execution.arguments:
-                    args = execution.arguments
-                    var_map = {}
-                    for key, val in args.items():
-                        part_list = val.part_list.parts
-                        if part_list and part_list[0].text:
-                            var_map[key] = part_list[0]
-                    if var_map:
-                        prompt.prompt_data.variables.append(var_map)
+        if api_schema.executions:
+            executions = api_schema.executions
+            if executions and prompt.prompt_data is not None:
+                prompt.prompt_data.variables = []
+                for execution in executions:
+                    if execution.arguments:
+                        args = execution.arguments
+                        var_map = {}
+                        for key, val in args.items():
+                            if (
+                                val.part_list is not None
+                                and val.part_list.parts is not None
+                            ):
+                                part_list = val.part_list.parts
+                                if part_list and part_list[0].text:
+                                    var_map[key] = part_list[0]
+                        if var_map and prompt.prompt_data.variables is not None:
+                            prompt.prompt_data.variables.append(var_map)
 
     return prompt
 

vertexai/_genai/_prompt_optimizer_utils.py

@@ -16,6 +16,7 @@
 
 import json
 from typing import Any, Optional, Union
+from typing_extensions import TypeAlias
 
 from pydantic import ValidationError
 
@@ -24,10 +25,10 @@
 try:
     import pandas as pd  # pylint: disable=g-import-not-at-top
 
-    PandasDataFrame = pd.DataFrame
+    PandasDataFrame: TypeAlias = pd.DataFrame
 except ImportError:
     pd = None
-    PandasDataFrame = Any
+    PandasDataFrame = Any  # type: ignore[misc]
 
 
 def _construct_input_prompt(
@@ -80,6 +81,12 @@ def _get_few_shot_prompt(
 ) -> str:
     """Builds the few shot prompt."""
 
+    if config.examples_dataframe is None:
+        raise ValueError("The 'examples_dataframe' is required in the config.")
+
+    if "prompt" not in config.examples_dataframe.columns:
+        raise ValueError("'prompt' is required in the examples_dataframe.")
+
     if "prompt" not in config.examples_dataframe.columns:
         raise ValueError("'prompt' is required in the examples_dataframe.")
     prompt_col_name = "prompt"
@@ -164,7 +171,7 @@ def _get_service_account(
         )
 
 
-def _clean_and_parse_optimized_prompt(output_str: str):
+def _clean_and_parse_optimized_prompt(output_str: str) -> Optional[Any]:
     """Cleans a string response returned from the prompt optimizer endpoint.
 
     Args:

vertexai/_genai/_transformers.py

@@ -14,7 +14,7 @@
 #
 
 """Transformers module for Vertex addons."""
-from typing import Any  # type: ignore[attr-defined]
+from typing import Any
 
 from google.genai._common import get_value_by_path as getv
 
@@ -75,7 +75,7 @@ def t_metrics(
                 pointwise_spec["system_instruction"] = system_instruction
             return_raw_output = getv(metric, ["return_raw_output"])
             if return_raw_output:
-                pointwise_spec["custom_output_format_config"] = {  # type: ignore[assignment]
+                pointwise_spec["custom_output_format_config"] = {
                     "return_raw_output": return_raw_output
                 }
             metric_payload_item["pointwise_metric_spec"] = pointwise_spec

vertexai/_genai/client.py

@@ -92,7 +92,7 @@ def evals(self) -> "evals_module.AsyncEvals":
                     "Please install them using pip install "
                     "google-cloud-aiplatform[evaluation]"
                 ) from e
-        return self._evals.AsyncEvals(self._api_client)
+        return self._evals.AsyncEvals(self._api_client)  # type: ignore[no-any-return]
 
     @property
     @_common.experimental_warning(
@@ -104,7 +104,7 @@ def prompt_optimizer(self) -> "prompt_optimizer_module.AsyncPromptOptimizer":
             self._prompt_optimizer = importlib.import_module(
                 ".prompt_optimizer", __package__
             )
-        return self._prompt_optimizer.AsyncPromptOptimizer(self._api_client)
+        return self._prompt_optimizer.AsyncPromptOptimizer(self._api_client)  # type: ignore[no-any-return]
 
     @property
     def agent_engines(self) -> "agent_engines_module.AsyncAgentEngines":
@@ -122,7 +122,7 @@ def agent_engines(self) -> "agent_engines_module.AsyncAgentEngines":
                     "Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._agent_engines.AsyncAgentEngines(self._api_client)
+        return self._agent_engines.AsyncAgentEngines(self._api_client)  # type: ignore[no-any-return]
 
     @property
     def prompts(self) -> "prompts_module.AsyncPrompts":
@@ -131,7 +131,7 @@ def prompts(self) -> "prompts_module.AsyncPrompts":
                 ".prompts",
                 __package__,
             )
-        return self._prompts.AsyncPrompts(self._api_client)
+        return self._prompts.AsyncPrompts(self._api_client)  # type: ignore[no-any-return]
 
     @property
     @_common.experimental_warning(
@@ -144,7 +144,7 @@ def datasets(self) -> "datasets_module.AsyncDatasets":
                 ".datasets",
                 __package__,
             )
-        return self._datasets.AsyncDatasets(self._api_client)
+        return self._datasets.AsyncDatasets(self._api_client)  # type: ignore[no-any-return]
 
     async def aclose(self) -> None:
         """Closes the async client explicitly.
@@ -255,7 +255,7 @@ def evals(self) -> "evals_module.Evals":
                     "Please install them using pip install "
                     "google-cloud-aiplatform[evaluation]"
                 ) from e
-        return self._evals.Evals(self._api_client)
+        return self._evals.Evals(self._api_client)  # type: ignore[no-any-return]
 
     @property
     @_common.experimental_warning(
@@ -267,7 +267,7 @@ def prompt_optimizer(self) -> "prompt_optimizer_module.PromptOptimizer":
             self._prompt_optimizer = importlib.import_module(
                 ".prompt_optimizer", __package__
             )
-        return self._prompt_optimizer.PromptOptimizer(self._api_client)
+        return self._prompt_optimizer.PromptOptimizer(self._api_client)  # type: ignore[no-any-return]
 
     @property
     def aio(self) -> "AsyncClient":
@@ -317,7 +317,7 @@ def agent_engines(self) -> "agent_engines_module.AgentEngines":
                     "Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._agent_engines.AgentEngines(self._api_client)
+        return self._agent_engines.AgentEngines(self._api_client)  # type: ignore[no-any-return]
 
     @property
     def prompts(self) -> "prompts_module.Prompts":
@@ -327,7 +327,7 @@ def prompts(self) -> "prompts_module.Prompts":
                 ".prompts",
                 __package__,
             )
-        return self._prompts.Prompts(self._api_client)
+        return self._prompts.Prompts(self._api_client)  # type: ignore[no-any-return]
 
     @property
     @_common.experimental_warning(
@@ -340,4 +340,4 @@ def datasets(self) -> "datasets_module.Datasets":
                 ".datasets",
                 __package__,
             )
-        return self._datasets.Datasets(self._api_client)
+        return self._datasets.Datasets(self._api_client)  # type: ignore[no-any-return]

vertexai/_genai/evals.py

@@ -1674,7 +1674,7 @@ def create_evaluation_run(
         if not name:
             name = f"evaluation_run_{uuid.uuid4()}"
 
-        return self._create_evaluation_run(  # type: ignore[no-any-return]
+        return self._create_evaluation_run(
             name=name,
             display_name=display_name or name,
             data_source=dataset,
@@ -1782,7 +1782,7 @@ def create_evaluation_item(
         Returns:
           The evaluation item.
         """
-        return self._create_evaluation_item(  # type: ignore[no-any-return]
+        return self._create_evaluation_item(
             evaluation_item_type=evaluation_item_type,
             gcs_uri=gcs_uri,
             display_name=display_name,
@@ -1812,7 +1812,7 @@ def create_evaluation_set(
         Returns:
           The evaluation set.
         """
-        return self._create_evaluation_set(  # type: ignore[no-any-return]
+        return self._create_evaluation_set(
             evaluation_items=evaluation_items,
             display_name=display_name,
             config=config,