mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Implement in memory EvalSetsManager
This version of the EvalSetsManager is intended to support two main behaviors 1) The agent developer wants to bring in their own eval set file, which is usually the case with `adk eval` cli. Once their eval sets are uploaded into this version of the eval sets manager, the EvalSetManager could be handed over to the Eval system for running evals. 2) As a part of AgentEvaluator testing, we expect developers to supply Eval cases in json files. The in-memory version of the EvalSetsManager will help us run those test cases using LocalEvalService. PiperOrigin-RevId: 783198788
This commit is contained in:
committed by
Copybara-Service
parent
a50419961c
commit
c25911c912
@@ -0,0 +1,151 @@
|
||||
# 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 typing import Optional
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from ..errors.not_found_error import NotFoundError
|
||||
from .eval_case import EvalCase
|
||||
from .eval_set import EvalSet
|
||||
from .eval_sets_manager import EvalSetsManager
|
||||
|
||||
|
||||
class InMemoryEvalSetsManager(EvalSetsManager):
|
||||
"""An in-memory implementation of EvalSetsManager using dictionaries.
|
||||
|
||||
You can use this class:
|
||||
1) As a part of your testcase.
|
||||
2) For cases where other implementations of EvalSetsManager are too expensive
|
||||
to use.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# {app_name: {eval_set_id: EvalSet}}
|
||||
self._eval_sets: dict[str, dict[str, EvalSet]] = {}
|
||||
# {app_name: {eval_set_id: {eval_case_id: EvalCase}}}
|
||||
self._eval_cases: dict[str, dict[str, dict[str, EvalCase]]] = {}
|
||||
|
||||
def _ensure_app_exists(self, app_name: str):
|
||||
if app_name not in self._eval_sets:
|
||||
self._eval_sets[app_name] = {}
|
||||
self._eval_cases[app_name] = {}
|
||||
|
||||
@override
|
||||
def get_eval_set(self, app_name: str, eval_set_id: str) -> Optional[EvalSet]:
|
||||
self._ensure_app_exists(app_name)
|
||||
return self._eval_sets[app_name].get(eval_set_id, None)
|
||||
|
||||
@override
|
||||
def create_eval_set(self, app_name: str, eval_set_id: str):
|
||||
self._ensure_app_exists(app_name)
|
||||
if eval_set_id in self._eval_sets[app_name]:
|
||||
raise ValueError(
|
||||
f"EvalSet {eval_set_id} already exists for app {app_name}."
|
||||
)
|
||||
|
||||
new_eval_set = EvalSet(
|
||||
eval_set_id=eval_set_id,
|
||||
eval_cases=[],
|
||||
creation_timestamp=time.time(),
|
||||
)
|
||||
self._eval_sets[app_name][eval_set_id] = new_eval_set
|
||||
self._eval_cases[app_name][eval_set_id] = {}
|
||||
|
||||
@override
|
||||
def list_eval_sets(self, app_name: str) -> list[str]:
|
||||
if app_name not in self._eval_sets:
|
||||
return []
|
||||
|
||||
return list(self._eval_sets[app_name].keys())
|
||||
|
||||
@override
|
||||
def get_eval_case(
|
||||
self, app_name: str, eval_set_id: str, eval_case_id: str
|
||||
) -> Optional[EvalCase]:
|
||||
if app_name not in self._eval_cases:
|
||||
return None
|
||||
if eval_set_id not in self._eval_cases[app_name]:
|
||||
return None
|
||||
return self._eval_cases[app_name][eval_set_id].get(eval_case_id)
|
||||
|
||||
@override
|
||||
def add_eval_case(self, app_name: str, eval_set_id: str, eval_case: EvalCase):
|
||||
self._ensure_app_exists(app_name)
|
||||
if eval_set_id not in self._eval_sets[app_name]:
|
||||
raise NotFoundError(
|
||||
f"EvalSet {eval_set_id} not found for app {app_name}."
|
||||
)
|
||||
if eval_case.eval_id in self._eval_cases[app_name][eval_set_id]:
|
||||
raise ValueError(
|
||||
f"EvalCase {eval_case.eval_id} already exists in EvalSet"
|
||||
f" {eval_set_id} for app {app_name}."
|
||||
)
|
||||
|
||||
self._eval_cases[app_name][eval_set_id][eval_case.eval_id] = eval_case
|
||||
# Also update the list in the EvalSet object
|
||||
self._eval_sets[app_name][eval_set_id].eval_cases.append(eval_case)
|
||||
|
||||
@override
|
||||
def update_eval_case(
|
||||
self, app_name: str, eval_set_id: str, updated_eval_case: EvalCase
|
||||
):
|
||||
self._ensure_app_exists(app_name)
|
||||
if eval_set_id not in self._eval_sets[app_name]:
|
||||
raise NotFoundError(
|
||||
f"EvalSet {eval_set_id} not found for app {app_name}."
|
||||
)
|
||||
if updated_eval_case.eval_id not in self._eval_cases[app_name][eval_set_id]:
|
||||
raise NotFoundError(
|
||||
f"EvalCase {updated_eval_case.eval_id} not found in EvalSet"
|
||||
f" {eval_set_id} for app {app_name}."
|
||||
)
|
||||
|
||||
# Full replace
|
||||
self._eval_cases[app_name][eval_set_id][
|
||||
updated_eval_case.eval_id
|
||||
] = updated_eval_case
|
||||
|
||||
# Update the list in the EvalSet object
|
||||
eval_set = self._eval_sets[app_name][eval_set_id]
|
||||
for i, case in enumerate(eval_set.eval_cases):
|
||||
if case.eval_id == updated_eval_case.eval_id:
|
||||
eval_set.eval_cases[i] = updated_eval_case
|
||||
break
|
||||
|
||||
@override
|
||||
def delete_eval_case(
|
||||
self, app_name: str, eval_set_id: str, eval_case_id: str
|
||||
):
|
||||
self._ensure_app_exists(app_name)
|
||||
if eval_set_id not in self._eval_sets[app_name]:
|
||||
raise NotFoundError(
|
||||
f"EvalSet {eval_set_id} not found for app {app_name}."
|
||||
)
|
||||
if eval_case_id not in self._eval_cases[app_name][eval_set_id]:
|
||||
raise NotFoundError(
|
||||
f"EvalCase {eval_case_id} not found in EvalSet {eval_set_id}"
|
||||
f" for app {app_name}."
|
||||
)
|
||||
|
||||
del self._eval_cases[app_name][eval_set_id][eval_case_id]
|
||||
|
||||
# Remove from the list in the EvalSet object
|
||||
eval_set = self._eval_sets[app_name][eval_set_id]
|
||||
eval_set.eval_cases = [
|
||||
case for case in eval_set.eval_cases if case.eval_id != eval_case_id
|
||||
]
|
||||
@@ -0,0 +1,199 @@
|
||||
# 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.
|
||||
|
||||
import time
|
||||
|
||||
from google.adk.errors.not_found_error import NotFoundError
|
||||
from google.adk.evaluation.eval_case import EvalCase
|
||||
from google.adk.evaluation.in_memory_eval_sets_manager import InMemoryEvalSetsManager
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_name():
|
||||
return "test_app"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def manager():
|
||||
return InMemoryEvalSetsManager()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def eval_set_id():
|
||||
return "test_eval_set"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def eval_case_id():
|
||||
return "test_eval_case"
|
||||
|
||||
|
||||
def test_create_eval_set(manager, app_name, eval_set_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_set = manager.get_eval_set(app_name, eval_set_id)
|
||||
assert eval_set is not None
|
||||
assert eval_set.eval_set_id == eval_set_id
|
||||
assert eval_set.eval_cases == []
|
||||
|
||||
|
||||
def test_create_eval_set_already_exists(manager, app_name, eval_set_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
with pytest.raises(ValueError):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
|
||||
|
||||
def test_get_eval_set(manager, app_name, eval_set_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_set = manager.get_eval_set(app_name, eval_set_id)
|
||||
assert eval_set is not None
|
||||
assert eval_set.eval_set_id == eval_set_id
|
||||
|
||||
|
||||
def test_get_eval_set_not_found(manager, app_name):
|
||||
eval_set = manager.get_eval_set(app_name, "nonexistent_set")
|
||||
assert eval_set is None
|
||||
|
||||
|
||||
def test_get_eval_set_wrong_app(manager, app_name, eval_set_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_set = manager.get_eval_set("wrong_app", eval_set_id)
|
||||
assert eval_set is None
|
||||
|
||||
|
||||
def test_list_eval_sets(manager, app_name):
|
||||
manager.create_eval_set(app_name, "set1")
|
||||
manager.create_eval_set(app_name, "set2")
|
||||
eval_sets = manager.list_eval_sets(app_name)
|
||||
assert len(eval_sets) == 2
|
||||
assert "set1" in eval_sets
|
||||
assert "set2" in eval_sets
|
||||
|
||||
|
||||
def test_list_eval_sets_wrong_app(manager, app_name):
|
||||
manager.create_eval_set(app_name, "set1")
|
||||
eval_sets = manager.list_eval_sets("wrong_app")
|
||||
assert len(eval_sets) == 0
|
||||
|
||||
|
||||
def test_add_eval_case(manager, app_name, eval_set_id, eval_case_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
|
||||
retrieved_case = manager.get_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
assert retrieved_case is not None
|
||||
assert retrieved_case.eval_id == eval_case_id
|
||||
|
||||
eval_set = manager.get_eval_set(app_name, eval_set_id)
|
||||
assert len(eval_set.eval_cases) == 1
|
||||
assert eval_set.eval_cases[0].eval_id == eval_case_id
|
||||
|
||||
|
||||
def test_add_eval_case_set_not_found(
|
||||
manager, app_name, eval_set_id, eval_case_id
|
||||
):
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
with pytest.raises(NotFoundError):
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
|
||||
|
||||
def test_add_eval_case_already_exists(
|
||||
manager, app_name, eval_set_id, eval_case_id
|
||||
):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
with pytest.raises(ValueError):
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
|
||||
|
||||
def test_get_eval_case(manager, app_name, eval_set_id, eval_case_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
retrieved_case = manager.get_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
assert retrieved_case is not None
|
||||
assert retrieved_case.eval_id == eval_case_id
|
||||
|
||||
|
||||
def test_get_eval_case_not_found(manager, app_name, eval_set_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
retrieved_case = manager.get_eval_case(
|
||||
app_name, eval_set_id, "nonexistent_case"
|
||||
)
|
||||
assert retrieved_case is None
|
||||
|
||||
|
||||
def test_get_eval_case_set_not_found(manager, app_name, eval_case_id):
|
||||
retrieved_case = manager.get_eval_case(
|
||||
app_name, "nonexistent_set", eval_case_id
|
||||
)
|
||||
assert retrieved_case is None
|
||||
|
||||
|
||||
def test_update_eval_case(manager, app_name, eval_set_id, eval_case_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
|
||||
updated_eval_case = EvalCase(
|
||||
eval_id=eval_case_id, conversation=[], creation_timestamp=time.time()
|
||||
)
|
||||
manager.update_eval_case(app_name, eval_set_id, updated_eval_case)
|
||||
|
||||
retrieved_case = manager.get_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
assert retrieved_case is not None
|
||||
assert retrieved_case.creation_timestamp != 0.0
|
||||
assert (
|
||||
retrieved_case.creation_timestamp == updated_eval_case.creation_timestamp
|
||||
)
|
||||
|
||||
eval_set = manager.get_eval_set(app_name, eval_set_id)
|
||||
assert len(eval_set.eval_cases) == 1
|
||||
assert (
|
||||
eval_set.eval_cases[0].creation_timestamp
|
||||
== updated_eval_case.creation_timestamp
|
||||
)
|
||||
|
||||
|
||||
def test_update_eval_case_not_found(
|
||||
manager, app_name, eval_set_id, eval_case_id
|
||||
):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
updated_eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
with pytest.raises(NotFoundError):
|
||||
manager.update_eval_case(app_name, eval_set_id, updated_eval_case)
|
||||
|
||||
|
||||
def test_delete_eval_case(manager, app_name, eval_set_id, eval_case_id):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
eval_case = EvalCase(eval_id=eval_case_id, conversation=[])
|
||||
manager.add_eval_case(app_name, eval_set_id, eval_case)
|
||||
|
||||
manager.delete_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
|
||||
retrieved_case = manager.get_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
assert retrieved_case is None
|
||||
|
||||
eval_set = manager.get_eval_set(app_name, eval_set_id)
|
||||
assert len(eval_set.eval_cases) == 0
|
||||
|
||||
|
||||
def test_delete_eval_case_not_found(
|
||||
manager, app_name, eval_set_id, eval_case_id
|
||||
):
|
||||
manager.create_eval_set(app_name, eval_set_id)
|
||||
with pytest.raises(NotFoundError):
|
||||
manager.delete_eval_case(app_name, eval_set_id, eval_case_id)
|
||||
Reference in New Issue
Block a user