8000 feat: Re-factor some eval sets manager logic, and implement GcsEvalSe… · ecda909/adk-python@1551bd4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 1551bd4

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: Re-factor some eval sets manager logic, and implement GcsEvalSetsManager to handle storage of eval sets on GCS
Eval sets will be stored as json files under `gs://{bucket_name}/{app_name}/evals/eval_sets/` PiperOrigin-RevId: 770487129
1 parent bbceb4f commit 1551bd4

File tree

5 files changed

+916
-98
lines changed

5 files changed

+916
-98
lines changed

‎src/google/adk/evaluation/_eval_sets_manager_utils.py

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import logging
18+
from typing import Optional
19+
20+
from ..errors.not_found_error import NotFoundError
21+
from .eval_case import EvalCase
22+
from .eval_set import EvalSet
23+
from .eval_sets_manager import EvalSetsManager
24+
25+
logger = logging.getLogger("google_adk." + __name__)
26+
27+
28+
def get_eval_set_from_app_and_id(
29+
eval_sets_manager: EvalSetsManager, app_name: str, eval_set_id: str
30+
) -> EvalSet:
31+
"""Returns an EvalSet if found, otherwise raises NotFoundError."""
32+
eval_set = eval_sets_manager.get_eval_set(app_name, eval_set_id)
33+
if not eval_set:
34+
raise NotFoundError(f"Eval set `{eval_set_id}` not found.")
35+
return eval_set
36+
37+
38+
def get_eval_case_from_eval_set(
39+
eval_set: EvalSet, eval_case_id: str
40+
) -> Optional[EvalCase]:
41+
"""Returns an EvalCase if found, otherwise None."""
42+
eval_case_to_find = None
43+
44+
# Look up the eval case by eval_case_id
45+
for eval_case in eval_set.eval_cases:
46+
if eval_case.eval_id == eval_case_id:
47+
eval_case_to_find = eval_case
48+
break
49+
50+
return eval_case_to_find
51+
52+
53+
def add_eval_case_to_eval_set(
54+
eval_set: EvalSet, eval_case: EvalCase
55+
) -> EvalSet:
56+
"""Adds an eval case to an eval set and returns the updated eval set."""
57+
eval_case_id = eval_case.eval_id
58+
59+
if [x for x in eval_set.eval_cases if x.eval_id == eval_case_id]:
60+
raise ValueError(
61+
f"Eval id `{eval_case_id}` already exists in `{eval_set.eval_set_id}`"
62+
" eval set.",
63+
)
64+
65+
eval_set.eval_cases.append(eval_case)
66+
return eval_set
67+
68+
69+
def update_eval_case_in_eval_set(
70+
eval_set: EvalSet, updated_eval_case: EvalCase
71+
) -> EvalSet:
72+
"""Updates an eval case in an eval set and returns the updated eval set."""
73+
# Find the eval case to be updated.
74+
eval_case_id = updated_eval_case.eval_id
75+
eval_case_to_update = get_eval_case_from_eval_set(eval_set, eval_case_id)
76+
77+
if not eval_case_to_update:
78+
raise NotFoundError(
79+
f"Eval case `{eval_case_id}` not found in eval set"
80+
f" `{eval_set.eval_set_id}`."
81+
)
82+
83+
# Remove the existing eval case and add the updated eval case.
84+
eval_set.eval_cases.remove(eval_case_to_update)
85+
eval_set.eval_cases.append(updated_eval_case)
86+
return eval_set
87+
88+
89+
def delete_eval_case_from_eval_set(
90+
eval_set: EvalSet, eval_case_id: str
91+
) -> EvalSet:
92+
"""Deletes an eval case from an eval set and returns the updated eval set."""
93+
# Find the eval case to be deleted.
94+
eval_case_to_delete = get_eval_case_from_eval_set(eval_set, eval_case_id)
95+
96+
if not eval_case_to_delete:
97+
raise NotFoundError(
98+
f"Eval case `{eval_case_id}` not found in eval set"
99+
f" `{eval_set.eval_set_id}`."
100+
)
101+
102+
# Remove the existing eval case.
103+
logger.info(
104+
"EvalCase`%s` was found in the eval set. It will be removed permanently.",
105+
eval_case_id,
106+
)
107+
eval_set.eval_cases.remove(eval_case_to_delete)
108+
return eval_set

‎src/google/adk/evaluation/gcs_eval_sets_manager.py

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import logging
18+
import re
19+
import time
20+
from typing import Optional
21+
22+
from google.cloud import exceptions as cloud_exceptions
23+
from google.cloud import storage
24+
from typing_extensions import override
25+
26+
from ._eval_sets_manager_utils import add_eval_case_to_eval_set
27+
from ._eval_sets_manager_utils import delete_eval_case_from_eval_set
28+
from ._eval_sets_manager_utils import get_eval_case_from_eval_set
29+
from ._eval_sets_manager_utils import get_eval_set_from_app_and_id
30+
from ._eval_sets_manager_utils import update_eval_case_in_eval_set
31+
from .eval_case import EvalCase
32+
from .eval_set import EvalSet
33+
from .eval_sets_manager import EvalSetsManager
34+
35+
logger = logging.getLogger("google_adk." + __name__)
36+
37+
_EVAL_SETS_DIR = "evals/eval_sets"
38+
_EVAL_SET_FILE_EXTENSION = ".evalset.json"
39+
40+
41+
class GcsEvalSetsManager(EvalSetsManager):
42+
"""An EvalSetsManager that stores eval sets in a GCS bucket."""
43+
44+
def __init__(self, bucket_name: str, **kwargs):
45+
"""Initializes the GcsEvalSetsManager.
46+
47+
Args:
48+
bucket_name: The name of the bucket to use.
49+
**kwargs: Keyword arguments to pass to the Google Cloud Storage client.
50+
"""
51+
self.bucket_name = bucket_name
52+
self.storage_client = storage.Client(**kwargs)
53+
self.bucket = self.storage_client.bucket(self.bucket_name)
54+
# Check if the bucket exists.
55+
if not self.bucket.exists():
56+
raise ValueError(
57+
f"Bucket `{self.bucket_name}` does not exist. Please create it "
58+
"before using the GcsEvalSetsManager."
59+
)
60+
61+
def _get_eval_sets_dir(self, app_name: str) -> str:
62+
return f"{app_name}/{_EVAL_SETS_DIR}"
63+
64+
def _get_eval_set_blob_name(self, app_name: str, eval_set_id: str) -> str:
65+
eval_sets_dir = self._get_eval_sets_dir(app_name)
66+
return f"{eval_sets_dir}/{eval_set_id}{_EVAL_SET_FILE_EXTENSION}"
67+
68+
def _validate_id(self, id_name: str, id_value: str):
69+
pattern = r"^[a-zA-Z0-9_]+$"
70+
if not bool(re.fullmatch(pattern, id_value)):
71+
raise ValueError(
72+
f"Invalid {id_name}. {id_name} should have the `{pattern}` format",
73+
)
74+
75+
def _write_eval_set_to_blob(self, blob_name: str, eval_set: EvalSet):
76+
"""Writes an EvalSet to GCS."""
77+
blob = self.bucket.blob(blob_name)
78+
blob.upload_from_string(
79+
eval_set.model_dump_json(indent=2),
80+
content_type="application/json",
81+
)
82+
83+
def _save_eval_set(self, app_name: str, eval_set_id: str, eval_set: EvalSet):
84+
eval_set_blob_name = self._get_eval_set_blob_name(app_name, eval_set_id)
85+
self._write_eval_set_to_blob(eval_set_blob_name, eval_set)
86+
87+
@override
88+
def get_eval_set(self, app_name: str, eval_set_id: str) -> Optional[EvalSet]:
89+
"""Returns an EvalSet identified by an app_name and eval_set_id."""
90+
eval_set_blob_name = self._get_eval_set_blob_name(app_name, eval_set_id)
91+
blob = self.bucket.blob(eval_set_blob_name)
92+
if not blob.exists():
93+
return None
94+
eval_set_data = blob.download_as_text()
95+
return EvalSet.model_validate_json(eval_set_data)
96+
97+
@override
98+
def create_eval_set(self, app_name: str, eval_set_id: str):
99+
"""Creates an empty EvalSet and saves it to GCS."""
100+
self._validate_id(id_name="Eval Set Id", id_value=eval_set_id)
101+
new_eval_set_blob_name = self._get_eval_set_blob_name(app_name, eval_set_id)
102+
if self.bucket.blob(new_eval_set_blob_name).exists():
103+
raise ValueError(
104+
f"Eval set `{eval_set_id}` already exists for app `{app_name}`."
105+
)
106+
logger.info("Creating eval set blob: `%s`", new_eval_set_blob_name)
107+
new_eval_set = EvalSet(
108+
eval_set_id=eval_set_id,
109+
name=eval_set_id,
110+
eval_cases=[],
111+
creation_timestamp=time.time(),
112+
)
113+
self._write_eval_set_to_blob(new_eval_set_blob_name, new_eval_set)
114+
115+
@override
116+
def list_eval_sets(self, app_name: str) -> list[str]:
117+
"""Returns a list of EvalSet ids that belong to the given app_name."""
118+
eval_sets_dir = self._get_eval_sets_dir(app_name)
119+
eval_sets = []
120+
try:
121+
for blob in self.bucket.list_blobs(prefix=eval_sets_dir):
122+
if not blob.name.endswith(_EVAL_SET_FILE_EXTENSION):
123+
continue
124+
eval_set_id = blob.name.split("/")[-1].removesuffix(
125+
_EVAL_SET_FILE_EXTENSION
126+
)
127+
eval_sets.append(eval_set_id)
128+
return sorted(eval_sets)
129+
except cloud_exceptions.NotFound as e:
130+
raise ValueError(
131+
f"App `{app_name}` not found in GCS bucket `{self.bucket_name}`."
132+
) from e
133+
134+
@override
135+
def get_eval_case(
136+
self, app_name: str, eval_set_id: str, eval_case_id: str
137+
) -> Optional[EvalCase]:
138+
"""Returns an EvalCase identified by an app_name, eval_set_id and eval_case_id."""
139+
eval_set = self.get_eval_set(app_name, eval_set_id)
140+
if not eval_set:
141+
return None
142+
return get_eval_case_from_eval_set(eval_set, eval_case_id)
143+
144+
@override
145+
def add_eval_case(self, app_name: str, eval_set_id: str, eval_case: EvalCase):
146+
"""Adds the given EvalCase to an existing EvalSet.
147+
148+
Args:
149+
app_name: The name of the app.
150+
eval_set_id: The id of the eval set containing the eval case to update.
151+
eval_case: The EvalCase to add.
152+
153+
Raises:
154+
NotFoundError: If the eval set is not found.
155+
ValueError: If the eval case already exists in the eval set.
156+
"""
157+
eval_set = get_eval_set_from_app_and_id(self, app_name, eval_set_id)
158+
updated_eval_set = add_eval_case_to_eval_set(eval_set, eval_case)
159+
self._save_eval_set(app_name, eval_set_id, updated_eval_set)
160+
161+
@override
162+
def update_eval_case(
163+
self, app_name: str, eval_set_id: str, updated_eval_case: EvalCase
164+
):
165+
"""Updates an existing EvalCase.
166+
167+
Args:
168+
app_name: The name of the app.
169+
eval_set_id: The id of the eval set containing the eval case to update.
170+
updated_eval_case: The updated EvalCase. Overwrites the existing EvalCase
171+
using the eval_id field.
172+
173+
Raises:
174+
NotFoundError: If the eval set or the eval case is not found.
175+
"""
176+
eval_set = get_eval_set_from_app_and_id(self, app_name, eval_set_id)
177+
updated_eval_set = update_eval_case_in_eval_set(eval_set, updated_eval_case)
178+
self._save_eval_set(app_name, eval_set_id, updated_eval_set)
179+
180+
@override
181+
def delete_eval_case(
182+
self, app_name: str, eval_set_id: str, eval_case_id: str
183+
):
184+
"""Deletes the EvalCase with the given eval_case_id from the given EvalSet.
185+
186+
Args:
187+
app_name: The name of the app.
188+
eval_set_id: The id of the eval set containing the eval case to delete.
189+
eval_case_id: The id of the eval case to delete.
190+
191+
Raises:
192+
NotFoundError: If the eval set or the eval case to delete is not found.
193+
"""
194+
eval_set = get_eval_set_from_app_and_id(self, app_name, eval_set_id)
195+
updated_eval_set = delete_eval_case_from_eval_set(eval_set, eval_case_id)
196+
self._save_eval_set(app_name, eval_set_id, updated_eval_set)

0 commit comments

Comments
 (0)
0