From 4ee125a03856fdb9ed28245bf7f5917c2d9038db Mon Sep 17 00:00:00 2001 From: Keyur Joshi Date: Tue, 27 Jan 2026 12:55:58 -0800 Subject: [PATCH] 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 PiperOrigin-RevId: 861849691 --- src/google/adk/optimization/__init__.py | 13 +++ .../adk/optimization/agent_optimizer.py | 49 ++++++++++ src/google/adk/optimization/data_types.py | 90 +++++++++++++++++++ src/google/adk/optimization/sampler.py | 70 +++++++++++++++ 4 files changed, 222 insertions(+) create mode 100644 src/google/adk/optimization/__init__.py create mode 100644 src/google/adk/optimization/agent_optimizer.py create mode 100644 src/google/adk/optimization/data_types.py create mode 100644 src/google/adk/optimization/sampler.py diff --git a/src/google/adk/optimization/__init__.py b/src/google/adk/optimization/__init__.py new file mode 100644 index 00000000..58d482ea --- /dev/null +++ b/src/google/adk/optimization/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 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. diff --git a/src/google/adk/optimization/agent_optimizer.py b/src/google/adk/optimization/agent_optimizer.py new file mode 100644 index 00000000..df2a8dbe --- /dev/null +++ b/src/google/adk/optimization/agent_optimizer.py @@ -0,0 +1,49 @@ +# Copyright 2026 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 + +from abc import ABC +from abc import abstractmethod +from typing import Generic + +from ..agents.llm_agent import Agent +from .data_types import AgentWithScores +from .data_types import OptimizerResult +from .data_types import SamplingResult +from .sampler import Sampler + + +class AgentOptimizer(ABC, Generic[SamplingResult, AgentWithScores]): + """Base class for agent optimizers.""" + + @abstractmethod + async def optimize( + self, + initial_agent: Agent, + sampler: Sampler[SamplingResult], + ) -> OptimizerResult[AgentWithScores]: + """Runs the optimizer. + + Args: + initial_agent: The initial agent to be optimized. + sampler: The interface used to get training and validation example UIDs, + request agent evaluations, and get useful data for optimizing the agent. + + Returns: + The final result of the optimization process, containing the optimized + agent instances along with their corresponding scores on the validation + examples and any optimization metadata. + """ + ... diff --git a/src/google/adk/optimization/data_types.py b/src/google/adk/optimization/data_types.py new file mode 100644 index 00000000..bab9cf7a --- /dev/null +++ b/src/google/adk/optimization/data_types.py @@ -0,0 +1,90 @@ +# Copyright 2026 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 + +from typing import Any +from typing import Generic +from typing import Optional +from typing import TypeVar + +from google.adk.agents.llm_agent import Agent +from pydantic import BaseModel +from pydantic import Field + + +class BaseSamplingResult(BaseModel): + """Base class for evaluation results of the candidate agent on the batch of examples.""" + + scores: dict[str, float] = Field( + required=True, + description=( + "A map from example UID to the agent's overall score on that example." + " (higher is better)." + ), + ) + + +# SamplingResult: the per-component evaluation results for a batch of examples. +# Should at least include per-example scores and may also contain other data +# required for optimizing the agent (e.g., outputs, trajectories, and metrics). +SamplingResult = TypeVar("SamplingResult", bound=BaseSamplingResult) + + +class BaseAgentWithScores(BaseModel): + """An optimized agent with its scores. + + Optimizers may use the overall_score field and can return custom metrics by + sub-classing this class. + """ + + optimized_agent: Agent = Field( + required=True, + description="The optimized agent.", + ) + + overall_score: Optional[float] = Field( + default=None, + description="The overall score of the optimized agent.", + ) + + +AgentWithScores = TypeVar("AgentWithScores", bound=BaseAgentWithScores) + + +class OptimizerResult(BaseModel, Generic[AgentWithScores]): + """Base class for optimizer final results.""" + + optimized_agents: list[AgentWithScores] = Field( + required=True, + description=( + "A list of optimized agents which cannot be considered strictly" + " better than one another (see" + " https://en.wikipedia.org/wiki/Pareto_front), along with scores." + ), + ) + + +class UnstructuredSamplingResult(BaseSamplingResult): + """Evaluation result providing per-example unstructured evaluation data.""" + + data: Optional[dict[str, dict[str, Any]]] = Field( + default=None, + description=( + "A map from example UID to JSON-serializable evaluation data useful" + " for agent optimization. Recommended contents include inputs," + " trajectories, and metrics. Must be provided if requested by the" + " optimizer." + ), + ) diff --git a/src/google/adk/optimization/sampler.py b/src/google/adk/optimization/sampler.py new file mode 100644 index 00000000..0a0ff45d --- /dev/null +++ b/src/google/adk/optimization/sampler.py @@ -0,0 +1,70 @@ +# Copyright 2026 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 + +from abc import ABC +from abc import abstractmethod +from typing import Generic +from typing import Literal +from typing import Optional + +from ..agents.llm_agent import Agent +from .data_types import SamplingResult + + +class Sampler(ABC, Generic[SamplingResult]): + """Base class for agent optimizers to sample and score candidate agents. + + The developer must implement this interface for their evaluation service to + work with the optimizer. The optimizer will call the sample_and_score method + to get evaluation results for the candidate agent on the batch of examples. + """ + + @abstractmethod + def get_train_example_ids(self) -> list[str]: + """Returns the UIDs of examples to use for training the agent.""" + ... + + @abstractmethod + def get_validation_example_ids(self) -> list[str]: + """Returns the UIDs of examples to use for validating the optimized agent.""" + ... + + @abstractmethod + async def sample_and_score( + self, + candidate: Agent, + example_set: Literal["train", "validation"] = "validation", + batch: Optional[list[str]] = None, + capture_full_eval_data: bool = False, + ) -> SamplingResult: + """Evaluates the candidate agent on the batch of examples. + + Args: + candidate: The candidate agent to be evaluated. + example_set: The set of examples to evaluate the candidate agent on. + Possible values are "train" and "validation". + batch: List of UIDs of examples to evaluate the candidate agent on. If not + provided, all examples from the chosen set will be used. + capture_full_eval_data: If false, it is enough to only calculate the + scores for each example. If true, this method should also capture all + other data required for optimizing the agent (e.g., outputs, + trajectories, and tool calls). + + Returns: + The evaluation results, containing the scores for each example and (if + requested) other data required for optimization. + """ + ...