chore: Add a2a task result aggregator

PiperOrigin-RevId: 775975982
This commit is contained in:
Xiang (Sean) Zhou
2025-06-25 22:02:28 -07:00
committed by Copybara-Service
parent 77b869f5e3
commit 2f55de6ded
4 changed files with 434 additions and 0 deletions
+13
View File
@@ -0,0 +1,13 @@
# 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.
@@ -0,0 +1,71 @@
# 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
from a2a.server.events import Event
from a2a.types import Message
from a2a.types import TaskState
from a2a.types import TaskStatusUpdateEvent
from ...utils.feature_decorator import working_in_progress
@working_in_progress
class TaskResultAggregator:
"""Aggregates the task status updates and provides the final task state."""
def __init__(self):
self._task_state = TaskState.working
self._task_status_message = None
def process_event(self, event: Event):
"""Process an event from the agent run and detect signals about the task status.
Priority of task state:
- failed
- auth_required
- input_required
- working
"""
if isinstance(event, TaskStatusUpdateEvent):
if event.status.state == TaskState.failed:
self._task_state = TaskState.failed
self._task_status_message = event.status.message
elif (
event.status.state == TaskState.auth_required
and self._task_state != TaskState.failed
):
self._task_state = TaskState.auth_required
self._task_status_message = event.status.message
elif (
event.status.state == TaskState.input_required
and self._task_state
not in (TaskState.failed, TaskState.auth_required)
):
self._task_state = TaskState.input_required
self._task_status_message = event.status.message
# final state is already recorded and make sure the intermediate state is
# always working because other state may terminate the event aggregation
# in a2a request handler
elif self._task_state == TaskState.working:
self._task_status_message = event.status.message
event.status.state = TaskState.working
@property
def task_state(self) -> TaskState:
return self._task_state
@property
def task_status_message(self) -> Message | None:
return self._task_status_message