feat: Implement GcsEvalSetResultsManager to handle storage of eval sets on GCS, and refactor eval set results manager

Eval results will be stored as json files under `gs://{bucket_name}/{app_name}/evals/eval_history/`

PiperOrigin-RevId: 770499242
This commit is contained in:
Google Team Member
2025-06-11 23:41:54 -07:00
committed by Copybara-Service
parent 1551bd4f4d
commit 0a5cf45a75
9 changed files with 503 additions and 155 deletions
@@ -0,0 +1,44 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import time
from .eval_result import EvalCaseResult
from .eval_result import EvalSetResult
def _sanitize_eval_set_result_name(eval_set_result_name: str) -> str:
"""Sanitizes the eval set result name."""
return eval_set_result_name.replace("/", "_")
def create_eval_set_result(
app_name: str,
eval_set_id: str,
eval_case_results: list[EvalCaseResult],
) -> EvalSetResult:
"""Creates a new EvalSetResult given eval_case_results."""
timestamp = time.time()
eval_set_result_id = f"{app_name}_{eval_set_id}_{timestamp}"
eval_set_result_name = _sanitize_eval_set_result_name(eval_set_result_id)
eval_set_result = EvalSetResult(
eval_set_result_id=eval_set_result_id,
eval_set_result_name=eval_set_result_name,
eval_set_id=eval_set_id,
eval_case_results=eval_case_results,
creation_timestamp=timestamp,
)
return eval_set_result
+12 -7
View File
@@ -36,8 +36,9 @@ class EvalCaseResult(BaseModel):
populate_by_name=True,
)
eval_set_file: str = Field(
eval_set_file: Optional[str] = Field(
deprecated=True,
default=None,
description="This field is deprecated, use eval_set_id instead.",
)
eval_set_id: str = ""
@@ -49,11 +50,15 @@ class EvalCaseResult(BaseModel):
final_eval_status: EvalStatus
"""Final eval status for this eval case."""
eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]] = Field(
deprecated=True,
description=(
"This field is deprecated, use overall_eval_metric_results instead."
),
eval_metric_results: Optional[list[tuple[EvalMetric, EvalMetricResult]]] = (
Field(
deprecated=True,
default=None,
description=(
"This field is deprecated, use overall_eval_metric_results"
" instead."
),
)
)
overall_eval_metric_results: list[EvalMetricResult]
@@ -80,7 +85,7 @@ class EvalSetResult(BaseModel):
populate_by_name=True,
)
eval_set_result_id: str
eval_set_result_name: str
eval_set_result_name: Optional[str] = None
eval_set_id: str
eval_case_results: list[EvalCaseResult] = Field(default_factory=list)
creation_timestamp: float = 0.0
@@ -16,6 +16,7 @@ from __future__ import annotations
from abc import ABC
from abc import abstractmethod
from typing import Optional
from .eval_result import EvalCaseResult
from .eval_result import EvalSetResult
@@ -38,7 +39,11 @@ class EvalSetResultsManager(ABC):
def get_eval_set_result(
self, app_name: str, eval_set_result_id: str
) -> EvalSetResult:
"""Returns an EvalSetResult identified by app_name and eval_set_result_id."""
"""Returns the EvalSetResult from app_name and eval_set_result_id.
Raises:
NotFoundError: If the EvalSetResult is not found.
"""
raise NotImplementedError()
@abstractmethod
@@ -0,0 +1,121 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import logging
from google.cloud import exceptions as cloud_exceptions
from google.cloud import storage
from typing_extensions import override
from ..errors.not_found_error import NotFoundError
from ._eval_set_results_manager_utils import create_eval_set_result
from .eval_result import EvalCaseResult
from .eval_result import EvalSetResult
from .eval_set_results_manager import EvalSetResultsManager
logger = logging.getLogger("google_adk." + __name__)
_EVAL_HISTORY_DIR = "evals/eval_history"
_EVAL_SET_RESULT_FILE_EXTENSION = ".evalset_result.json"
class GcsEvalSetResultsManager(EvalSetResultsManager):
"""An EvalSetResultsManager that stores eval results in a GCS bucket."""
def __init__(self, bucket_name: str, **kwargs):
"""Initializes the GcsEvalSetsManager.
Args:
bucket_name: The name of the bucket to use.
**kwargs: Keyword arguments to pass to the Google Cloud Storage client.
"""
self.bucket_name = bucket_name
self.storage_client = storage.Client(**kwargs)
self.bucket = self.storage_client.bucket(self.bucket_name)
# Check if the bucket exists.
if not self.bucket.exists():
raise ValueError(
f"Bucket `{self.bucket_name}` does not exist. Please create it before"
" using the GcsEvalSetsManager."
)
def _get_eval_history_dir(self, app_name: str) -> str:
return f"{app_name}/{_EVAL_HISTORY_DIR}"
def _get_eval_set_result_blob_name(
self, app_name: str, eval_set_result_id: str
) -> str:
eval_history_dir = self._get_eval_history_dir(app_name)
return f"{eval_history_dir}/{eval_set_result_id}{_EVAL_SET_RESULT_FILE_EXTENSION}"
def _write_eval_set_result(
self, blob_name: str, eval_set_result: EvalSetResult
):
"""Writes an EvalSetResult to GCS."""
blob = self.bucket.blob(blob_name)
blob.upload_from_string(
eval_set_result.model_dump_json(indent=2),
content_type="application/json",
)
@override
def save_eval_set_result(
self,
app_name: str,
eval_set_id: str,
eval_case_results: list[EvalCaseResult],
) -> None:
"""Creates and saves a new EvalSetResult given eval_case_results."""
eval_set_result = create_eval_set_result(
app_name, eval_set_id, eval_case_results
)
eval_set_result_blob_name = self._get_eval_set_result_blob_name(
app_name, eval_set_result.eval_set_result_id
)
logger.info("Writing eval result to blob: %s", eval_set_result_blob_name)
self._write_eval_set_result(eval_set_result_blob_name, eval_set_result)
@override
def get_eval_set_result(
self, app_name: str, eval_set_result_id: str
) -> EvalSetResult:
"""Returns an EvalSetResult from app_name and eval_set_result_id."""
eval_set_result_blob_name = self._get_eval_set_result_blob_name(
app_name, eval_set_result_id
)
blob = self.bucket.blob(eval_set_result_blob_name)
if not blob.exists():
raise NotFoundError(f"Eval set result `{eval_set_result_id}` not found.")
eval_set_result_data = blob.download_as_text()
return EvalSetResult.model_validate_json(eval_set_result_data)
@override
def list_eval_set_results(self, app_name: str) -> list[str]:
"""Returns the eval result ids that belong to the given app_name."""
eval_history_dir = self._get_eval_history_dir(app_name)
eval_set_results = []
try:
for blob in self.bucket.list_blobs(prefix=eval_history_dir):
eval_set_result_id = blob.name.split("/")[-1].removesuffix(
_EVAL_SET_RESULT_FILE_EXTENSION
)
eval_set_results.append(eval_set_result_id)
return sorted(eval_set_results)
except cloud_exceptions.NotFound as e:
raise ValueError(
f"App `{app_name}` not found in GCS bucket `{self.bucket_name}`."
) from e
@@ -17,10 +17,11 @@ from __future__ import annotations
import json
import logging
import os
import time
from typing_extensions import override
from ..errors.not_found_error import NotFoundError
from ._eval_set_results_manager_utils import create_eval_set_result
from .eval_result import EvalCaseResult
from .eval_result import EvalSetResult
from .eval_set_results_manager import EvalSetResultsManager
@@ -31,10 +32,6 @@ _ADK_EVAL_HISTORY_DIR = ".adk/eval_history"
_EVAL_SET_RESULT_FILE_EXTENSION = ".evalset_result.json"
def _sanitize_eval_set_result_name(eval_set_result_name: str) -> str:
return eval_set_result_name.replace("/", "_")
class LocalEvalSetResultsManager(EvalSetResultsManager):
"""An EvalSetResult manager that stores eval set results locally on disk."""
@@ -49,15 +46,8 @@ class LocalEvalSetResultsManager(EvalSetResultsManager):
eval_case_results: list[EvalCaseResult],
) -> None:
"""Creates and saves a new EvalSetResult given eval_case_results."""
timestamp = time.time()
eval_set_result_id = app_name + "_" + eval_set_id + "_" + str(timestamp)
eval_set_result_name = _sanitize_eval_set_result_name(eval_set_result_id)
eval_set_result = EvalSetResult(
eval_set_result_id=eval_set_result_id,
eval_set_result_name=eval_set_result_name,
eval_set_id=eval_set_id,
eval_case_results=eval_case_results,
creation_timestamp=timestamp,
eval_set_result = create_eval_set_result(
app_name, eval_set_id, eval_case_results
)
# Write eval result file, with eval_set_result_name.
app_eval_history_dir = self._get_eval_history_dir(app_name)
@@ -67,7 +57,7 @@ class LocalEvalSetResultsManager(EvalSetResultsManager):
eval_set_result_json = eval_set_result.model_dump_json()
eval_set_result_file_path = os.path.join(
app_eval_history_dir,
eval_set_result_name + _EVAL_SET_RESULT_FILE_EXTENSION,
eval_set_result.eval_set_result_name + _EVAL_SET_RESULT_FILE_EXTENSION,
)
logger.info("Writing eval result to file: %s", eval_set_result_file_path)
with open(eval_set_result_file_path, "w") as f:
@@ -87,9 +77,7 @@ class LocalEvalSetResultsManager(EvalSetResultsManager):
+ _EVAL_SET_RESULT_FILE_EXTENSION
)
if not os.path.exists(maybe_eval_result_file_path):
raise ValueError(
f"Eval set result `{eval_set_result_id}` does not exist."
)
raise NotFoundError(f"Eval set result `{eval_set_result_id}` not found.")
with open(maybe_eval_result_file_path, "r") as file:
eval_result_data = json.load(file)
return EvalSetResult.model_validate_json(eval_result_data)