feat: Add a new option eval_storage_uri in adk web & adk eval to specify GCS bucket to store eval data

PiperOrigin-RevId: 774947795

src/google/adk/cli/cli_tools_click.py

@@ -31,12 +31,15 @@ import uvicorn
 from . import cli_create
 from . import cli_deploy
 from .. import version
+from ..evaluation.gcs_eval_set_results_manager import GcsEvalSetResultsManager
+from ..evaluation.gcs_eval_sets_manager import GcsEvalSetsManager
 from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
 from ..sessions.in_memory_session_service import InMemorySessionService
 from .cli import run_cli
 from .cli_eval import MISSING_EVAL_DEPENDENCIES_MESSAGE
 from .fast_api import get_fast_api_app
 from .utils import envs
+from .utils import evals
 from .utils import logs
 
 LOG_LEVELS = click.Choice(
@@ -282,11 +285,21 @@ def cli_run(
     default=False,
     help="Optional. Whether to print detailed results on console or not.",
 )
+@click.option(
+    "--eval_storage_uri",
+    type=str,
+    help=(
+        "Optional. The evals storage URI to store agent evals,"
+        " supported URIs: gs://<bucket name>."
+    ),
+    default=None,
+)
 def cli_eval(
     agent_module_file_path: str,
-    eval_set_file_path: tuple[str],
+    eval_set_file_path: list[str],
     config_file_path: str,
     print_detailed_results: bool,
+    eval_storage_uri: Optional[str] = None,
 ):
   """Evaluates an agent given the eval sets.
 
@@ -338,11 +351,32 @@ def cli_eval(
   root_agent = get_root_agent(agent_module_file_path)
   reset_func = try_get_reset_func(agent_module_file_path)
 
+  gcs_eval_sets_manager = None
+  eval_set_results_manager = None
+  if eval_storage_uri:
+    gcs_eval_managers = evals.create_gcs_eval_managers_from_uri(
+        eval_storage_uri
+    )
+    gcs_eval_sets_manager = gcs_eval_managers.eval_sets_manager
+    eval_set_results_manager = gcs_eval_managers.eval_set_results_manager
+  else:
+    eval_set_results_manager = LocalEvalSetResultsManager(
+        agents_dir=os.path.dirname(agent_module_file_path)
+    )
   eval_set_file_path_to_evals = parse_and_get_evals_to_run(eval_set_file_path)
   eval_set_id_to_eval_cases = {}
 
   # Read the eval_set files and get the cases.
   for eval_set_file_path, eval_case_ids in eval_set_file_path_to_evals.items():
+    if gcs_eval_sets_manager:
+      eval_set = gcs_eval_sets_manager._load_eval_set_from_blob(
+          eval_set_file_path
+      )
+      if not eval_set:
+        raise click.ClickException(
+            f"Eval set {eval_set_file_path} not found in GCS."
+        )
+    else:
       eval_set = load_eval_set_from_file(eval_set_file_path, eval_set_file_path)
     eval_cases = eval_set.eval_cases
 
@@ -378,16 +412,13 @@ def cli_eval(
     raise click.ClickException(MISSING_EVAL_DEPENDENCIES_MESSAGE)
 
   # Write eval set results.
-  local_eval_set_results_manager = LocalEvalSetResultsManager(
-      agents_dir=os.path.dirname(agent_module_file_path)
-  )
   eval_set_id_to_eval_results = collections.defaultdict(list)
   for eval_case_result in eval_results:
     eval_set_id = eval_case_result.eval_set_id
     eval_set_id_to_eval_results[eval_set_id].append(eval_case_result)
 
   for eval_set_id, eval_case_results in eval_set_id_to_eval_results.items():
-    local_eval_set_results_manager.save_eval_set_result(
+    eval_set_results_manager.save_eval_set_result(
         app_name=os.path.basename(agent_module_file_path),
         eval_set_id=eval_set_id,
         eval_case_results=eval_case_results,
@@ -444,6 +475,15 @@ def adk_services_options():
         ),
         default=None,
     )
+    @click.option(
+        "--eval_storage_uri",
+        type=str,
+        help=(
+            "Optional. The evals storage URI to store agent evals,"
+            " supported URIs: gs://<bucket name>."
+        ),
+        default=None,
+    )
     @click.option(
         "--memory_service_uri",
         type=str,
@@ -564,6 +604,7 @@ def fast_api_common_options():
 )
 def cli_web(
     agents_dir: str,
+    eval_storage_uri: Optional[str] = None,
     log_level: str = "INFO",
     allow_origins: Optional[list[str]] = None,
     host: str = "127.0.0.1",
@@ -616,6 +657,7 @@ def cli_web(
       session_service_uri=session_service_uri,
       artifact_service_uri=artifact_service_uri,
       memory_service_uri=memory_service_uri,
+      eval_storage_uri=eval_storage_uri,
       allow_origins=allow_origins,
       web=True,
       trace_to_cloud=trace_to_cloud,
@@ -654,6 +696,7 @@ def cli_web(
 )
 def cli_api_server(
     agents_dir: str,
+    eval_storage_uri: Optional[str] = None,
     log_level: str = "INFO",
     allow_origins: Optional[list[str]] = None,
     host: str = "127.0.0.1",
@@ -685,6 +728,7 @@ def cli_api_server(
           session_service_uri=session_service_uri,
           artifact_service_uri=artifact_service_uri,
           memory_service_uri=memory_service_uri,
+          eval_storage_uri=eval_storage_uri,
           allow_origins=allow_origins,
           web=False,
           trace_to_cloud=trace_to_cloud,
@@ -771,6 +815,15 @@ def cli_api_server(
         " version in the dev environment)"
     ),
 )
+@click.option(
+    "--eval_storage_uri",
+    type=str,
+    help=(
+        "Optional. The evals storage URI to store agent evals,"
+        " supported URIs: gs://<bucket name>."
+    ),
+    default=None,
+)
 @adk_services_options()
 @deprecated_adk_services_options()
 @click.argument(
@@ -797,6 +850,7 @@ def cli_deploy_cloud_run(
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
+    eval_storage_uri: Optional[str] = None,
     session_db_url: Optional[str] = None,  # Deprecated
     artifact_storage_uri: Optional[str] = None,  # Deprecated
 ):

src/google/adk/cli/fast_api.py

@@ -65,6 +65,8 @@ from ..evaluation.eval_metrics import EvalMetric
 from ..evaluation.eval_metrics import EvalMetricResult
 from ..evaluation.eval_metrics import EvalMetricResultPerInvocation
 from ..evaluation.eval_result import EvalSetResult
+from ..evaluation.gcs_eval_set_results_manager import GcsEvalSetResultsManager
+from ..evaluation.gcs_eval_sets_manager import GcsEvalSetsManager
 from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
 from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
 from ..events.event import Event
@@ -198,6 +200,7 @@ def get_fast_api_app(
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
+    eval_storage_uri: Optional[str] = None,
     allow_origins: Optional[list[str]] = None,
     web: bool,
     trace_to_cloud: bool = False,
@@ -256,6 +259,16 @@ def get_fast_api_app(
 
   runner_dict = {}
 
+  # Set up eval managers.
+  eval_sets_manager = None
+  eval_set_results_manager = None
+  if eval_storage_uri:
+    gcs_eval_managers = evals.create_gcs_eval_managers_from_uri(
+        eval_storage_uri
+    )
+    eval_sets_manager = gcs_eval_managers.eval_sets_manager
+    eval_set_results_manager = gcs_eval_managers.eval_set_results_manager
+  else:
     eval_sets_manager = LocalEvalSetsManager(agents_dir=agents_dir)
     eval_set_results_manager = LocalEvalSetResultsManager(agents_dir=agents_dir)
 

src/google/adk/cli/utils/evals.py

@@ -14,17 +14,36 @@
 
 from __future__ import annotations
 
+import dataclasses
+import os
 from typing import Any
 from typing import Tuple
 
 from google.genai import types as genai_types
+from pydantic import alias_generators
+from pydantic import BaseModel
+from pydantic import ConfigDict
 from typing_extensions import deprecated
 
 from ...evaluation.eval_case import IntermediateData
 from ...evaluation.eval_case import Invocation
+from ...evaluation.gcs_eval_set_results_manager import GcsEvalSetResultsManager
+from ...evaluation.gcs_eval_sets_manager import GcsEvalSetsManager
 from ...sessions.session import Session
 
 
+class GcsEvalManagers(BaseModel):
+  model_config = ConfigDict(
+      alias_generator=alias_generators.to_camel,
+      populate_by_name=True,
+      arbitrary_types_allowed=True,
+  )
+
+  eval_sets_manager: GcsEvalSetsManager
+
+  eval_set_results_manager: GcsEvalSetResultsManager
+
+
 @deprecated('Use convert_session_to_eval_invocations instead.')
 def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]:
   """Converts a session data into eval format.
@@ -176,3 +195,37 @@ def convert_session_to_eval_invocations(session: Session) -> list[Invocation]:
       )
 
   return invocations
+
+
+def create_gcs_eval_managers_from_uri(
+    eval_storage_uri: str,
+) -> GcsEvalManagers:
+  """Creates GcsEvalManagers from eval_storage_uri.
+
+  Args:
+      eval_storage_uri: The evals storage URI to use. Supported URIs:
+        gs://<bucket name>. If a path is provided, the bucket will be extracted.
+
+  Returns:
+      GcsEvalManagers: The GcsEvalManagers object.
+
+  Raises:
+      ValueError: If the eval_storage_uri is not supported.
+  """
+  if eval_storage_uri.startswith('gs://'):
+    gcs_bucket = eval_storage_uri.split('://')[1]
+    eval_sets_manager = GcsEvalSetsManager(
+        bucket_name=gcs_bucket, project=os.environ['GOOGLE_CLOUD_PROJECT']
+    )
+    eval_set_results_manager = GcsEvalSetResultsManager(
+        bucket_name=gcs_bucket, project=os.environ['GOOGLE_CLOUD_PROJECT']
+    )
+    return GcsEvalManagers(
+        eval_sets_manager=eval_sets_manager,
+        eval_set_results_manager=eval_set_results_manager,
+    )
+  else:
+    raise ValueError(
+        f'Unsupported evals storage URI: {eval_storage_uri}. Supported URIs:'
+        ' gs://<bucket name>'
+    )

src/google/adk/evaluation/gcs_eval_sets_manager.py

@@ -72,6 +72,13 @@ class GcsEvalSetsManager(EvalSetsManager):
           f"Invalid {id_name}. {id_name} should have the `{pattern}` format",
       )
 
+  def _load_eval_set_from_blob(self, blob_name: str) -> Optional[EvalSet]:
+    blob = self.bucket.blob(blob_name)
+    if not blob.exists():
+      return None
+    eval_set_data = blob.download_as_text()
+    return EvalSet.model_validate_json(eval_set_data)
+
   def _write_eval_set_to_blob(self, blob_name: str, eval_set: EvalSet):
     """Writes an EvalSet to GCS."""
     blob = self.bucket.blob(blob_name)
@@ -88,11 +95,7 @@ class GcsEvalSetsManager(EvalSetsManager):
   def get_eval_set(self, app_name: str, eval_set_id: str) -> Optional[EvalSet]:
     """Returns an EvalSet identified by an app_name and eval_set_id."""
     eval_set_blob_name = self._get_eval_set_blob_name(app_name, eval_set_id)
-    blob = self.bucket.blob(eval_set_blob_name)
-    if not blob.exists():
-      return None
-    eval_set_data = blob.download_as_text()
-    return EvalSet.model_validate_json(eval_set_data)
+    return self._load_eval_set_from_blob(eval_set_blob_name)
 
   @override
   def create_eval_set(self, app_name: str, eval_set_id: str):

tests/unittests/cli/test_fast_api.py

@@ -40,7 +40,7 @@ logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 )
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 
 # Here we create a dummy agent module that get_fast_api_app expects
@@ -138,6 +138,7 @@ async def mock_run_evals_for_fast_api(*args, **kwargs):
       final_eval_status=1,  # Matches expected (assuming 1 is PASSED)
       user_id="test_user",  # Placeholder, adapt if needed
       session_id="test_session_for_eval_case",  # Placeholder
+      eval_set_file="test_eval_set_file",  # Placeholder
       overall_eval_metric_results=[{  # Matches expected
           "metricName": "tool_trajectory_avg_score",
           "threshold": 0.5,
@@ -372,7 +373,7 @@ def mock_eval_sets_manager():
 
 @pytest.fixture
 def mock_eval_set_results_manager():
-  """Create a mock eval set results manager."""
+  """Create a mock local eval set results manager."""
 
   # Storage for eval set results.
   eval_set_results = {}