Skip to content

Commit 4ee125a

Browse files
KPJoshicopybara-github
authored andcommitted
feat: Add interface for agent optimizers
Details: * The Sampler allows ADK agent optimizers to request evals and receive detailed eval results that can be used to guide agent optimization. * This interface allows developers to run custom evaluations if needed for their agent. An implementation to interface with the inbuilt LocalEvalService shall be published shortly in a follow-up PR. * The AgentOptimizer interface describes the general framework for an ADK agent optimizer that uses the Sampler interface to optimize an ADK agent. * data_types.py contains generic types and base classes to allow intercommunication between the AgentOptimizer, Sampler, and developer code. Co-authored-by: Keyur Joshi <keyurj@google.com> PiperOrigin-RevId: 861849691
1 parent 2155a35 commit 4ee125a

File tree

4 files changed

+222
-0
lines changed

4 files changed

+222
-0
lines changed
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2026 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.
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright 2026 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+
from abc import ABC
18+
from abc import abstractmethod
19+
from typing import Generic
20+
21+
from ..agents.llm_agent import Agent
22+
from .data_types import AgentWithScores
23+
from .data_types import OptimizerResult
24+
from .data_types import SamplingResult
25+
from .sampler import Sampler
26+
27+
28+
class AgentOptimizer(ABC, Generic[SamplingResult, AgentWithScores]):
29+
"""Base class for agent optimizers."""
30+
31+
@abstractmethod
32+
async def optimize(
33+
self,
34+
initial_agent: Agent,
35+
sampler: Sampler[SamplingResult],
36+
) -> OptimizerResult[AgentWithScores]:
37+
"""Runs the optimizer.
38+
39+
Args:
40+
initial_agent: The initial agent to be optimized.
41+
sampler: The interface used to get training and validation example UIDs,
42+
request agent evaluations, and get useful data for optimizing the agent.
43+
44+
Returns:
45+
The final result of the optimization process, containing the optimized
46+
agent instances along with their corresponding scores on the validation
47+
examples and any optimization metadata.
48+
"""
49+
...
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Copyright 2026 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+
from typing import Any
18+
from typing import Generic
19+
from typing import Optional
20+
from typing import TypeVar
21+
22+
from google.adk.agents.llm_agent import Agent
23+
from pydantic import BaseModel
24+
from pydantic import Field
25+
26+
27+
class BaseSamplingResult(BaseModel):
28+
"""Base class for evaluation results of the candidate agent on the batch of examples."""
29+
30+
scores: dict[str, float] = Field(
31+
required=True,
32+
description=(
33+
"A map from example UID to the agent's overall score on that example."
34+
" (higher is better)."
35+
),
36+
)
37+
38+
39+
# SamplingResult: the per-component evaluation results for a batch of examples.
40+
# Should at least include per-example scores and may also contain other data
41+
# required for optimizing the agent (e.g., outputs, trajectories, and metrics).
42+
SamplingResult = TypeVar("SamplingResult", bound=BaseSamplingResult)
43+
44+
45+
class BaseAgentWithScores(BaseModel):
46+
"""An optimized agent with its scores.
47+
48+
Optimizers may use the overall_score field and can return custom metrics by
49+
sub-classing this class.
50+
"""
51+
52+
optimized_agent: Agent = Field(
53+
required=True,
54+
description="The optimized agent.",
55+
)
56+
57+
overall_score: Optional[float] = Field(
58+
default=None,
59+
description="The overall score of the optimized agent.",
60+
)
61+
62+
63+
AgentWithScores = TypeVar("AgentWithScores", bound=BaseAgentWithScores)
64+
65+
66+
class OptimizerResult(BaseModel, Generic[AgentWithScores]):
67+
"""Base class for optimizer final results."""
68+
69+
optimized_agents: list[AgentWithScores] = Field(
70+
required=True,
71+
description=(
72+
"A list of optimized agents which cannot be considered strictly"
73+
" better than one another (see"
74+
" https://en.wikipedia.org/wiki/Pareto_front), along with scores."
75+
),
76+
)
77+
78+
79+
class UnstructuredSamplingResult(BaseSamplingResult):
80+
"""Evaluation result providing per-example unstructured evaluation data."""
81+
82+
data: Optional[dict[str, dict[str, Any]]] = Field(
83+
default=None,
84+
description=(
85+
"A map from example UID to JSON-serializable evaluation data useful"
86+
" for agent optimization. Recommended contents include inputs,"
87+
" trajectories, and metrics. Must be provided if requested by the"
88+
" optimizer."
89+
),
90+
)
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright 2026 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+
from abc import ABC
18+
from abc import abstractmethod
19+
from typing import Generic
20+
from typing import Literal
21+
from typing import Optional
22+
23+
from ..agents.llm_agent import Agent
24+
from .data_types import SamplingResult
25+
26+
27+
class Sampler(ABC, Generic[SamplingResult]):
28+
"""Base class for agent optimizers to sample and score candidate agents.
29+
30+
The developer must implement this interface for their evaluation service to
31+
work with the optimizer. The optimizer will call the sample_and_score method
32+
to get evaluation results for the candidate agent on the batch of examples.
33+
"""
34+
35+
@abstractmethod
36+
def get_train_example_ids(self) -> list[str]:
37+
"""Returns the UIDs of examples to use for training the agent."""
38+
...
39+
40+
@abstractmethod
41+
def get_validation_example_ids(self) -> list[str]:
42+
"""Returns the UIDs of examples to use for validating the optimized agent."""
43+
...
44+
45+
@abstractmethod
46+
async def sample_and_score(
47+
self,
48+
candidate: Agent,
49+
example_set: Literal["train", "validation"] = "validation",
50+
batch: Optional[list[str]] = None,
51+
capture_full_eval_data: bool = False,
52+
) -> SamplingResult:
53+
"""Evaluates the candidate agent on the batch of examples.
54+
55+
Args:
56+
candidate: The candidate agent to be evaluated.
57+
example_set: The set of examples to evaluate the candidate agent on.
58+
Possible values are "train" and "validation".
59+
batch: List of UIDs of examples to evaluate the candidate agent on. If not
60+
provided, all examples from the chosen set will be used.
61+
capture_full_eval_data: If false, it is enough to only calculate the
62+
scores for each example. If true, this method should also capture all
63+
other data required for optimizing the agent (e.g., outputs,
64+
trajectories, and tool calls).
65+
66+
Returns:
67+
The evaluation results, containing the scores for each example and (if
68+
requested) other data required for optimization.
69+
"""
70+
...

0 commit comments

Comments
 (0)