Skip to content

Commit 8cd3b12

Browse files
yang0228i-yliu
authored and
copybara-github
committed
fix(eval): preserve live usage without extra inference calls
Merge #7323 Gemini Live may send token usage separately from response content, and evaluation conversion dropped those reports. Preserve usage-bearing events, including empty and non-text content, and merge standalone usage into a compatible model event (same author, model and Live connection) without mutating input events or adding inference calls. A report with no compatible model event is kept and counts as the one model call it represents. Fixes #7321 Co-authored-by: Yi Liu <[email protected]> COPYBARA_INTEGRATE_REVIEW=#7323 from yang0228:codex/7321-live-eval-usage 2d2d71d PiperOrigin-RevId: 1011019927
1 parent af371b7 commit 8cd3b12

2 files changed

Lines changed: 380 additions & 7 deletions

File tree

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

Lines changed: 59 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1055,7 +1055,10 @@ def convert_events_to_eval_invocations(
10551055
final_response = event.content
10561056
final_event = event
10571057

1058-
should_add_event = event.grounding_metadata is not None
1058+
should_add_event = (
1059+
event.grounding_metadata is not None
1060+
or event.usage_metadata is not None
1061+
)
10591062
for p in event.content.parts:
10601063
if (
10611064
p.function_call
@@ -1067,12 +1070,61 @@ def convert_events_to_eval_invocations(
10671070
break
10681071
if should_add_event:
10691072
events_to_add.append(event)
1070-
elif event.grounding_metadata is not None:
1073+
elif (
1074+
event.grounding_metadata is not None
1075+
or event.usage_metadata is not None
1076+
):
10711077
events_to_add.append(event)
10721078

1073-
invocation_events = []
1074-
for e in events_to_add:
1075-
invocation_events.append(
1079+
# Pair each standalone usage report with a model event before building
1080+
# any `InvocationEvent`, so every record is constructed once with its
1081+
# final usage. Merging keeps the tokens without the report counting as
1082+
# a second model call. Nothing in the payload identifies the originating
1083+
# call, so the nearest unclaimed compatible event wins, earlier events
1084+
# first. Totals and call counts are exact; which event carries a report
1085+
# is best effort (e.g. `GeminiLlmConnection.receive` yields a message's
1086+
# usage before its content, so it can land on an earlier call).
1087+
# TODO: Index by author, model and Live session if long invocations
1088+
# make this quadratic search expensive.
1089+
merged_usage: dict[int, types.GenerateContentResponseUsageMetadata] = {}
1090+
merged_report_indices: set[int] = set()
1091+
for index, event in enumerate(events_to_add):
1092+
if (
1093+
(event.content and event.content.parts)
1094+
or event.grounding_metadata is not None
1095+
or event.usage_metadata is None
1096+
):
1097+
continue
1098+
target_index = next(
1099+
(
1100+
candidate_index
1101+
for candidate_index in (
1102+
*range(index - 1, -1, -1),
1103+
*range(index + 1, len(events_to_add)),
1104+
)
1105+
if candidate_index not in merged_usage
1106+
and events_to_add[candidate_index].usage_metadata is None
1107+
and events_to_add[candidate_index].model_version is not None
1108+
and events_to_add[candidate_index].author == event.author
1109+
and events_to_add[candidate_index].live_session_id
1110+
== event.live_session_id
1111+
and (
1112+
event.model_version is None
1113+
or events_to_add[candidate_index].model_version
1114+
== event.model_version
1115+
)
1116+
),
1117+
None,
1118+
)
1119+
if target_index is not None:
1120+
merged_usage[target_index] = event.usage_metadata
1121+
merged_report_indices.add(index)
1122+
1123+
merged_events = []
1124+
for index, e in enumerate(events_to_add):
1125+
if index in merged_report_indices:
1126+
continue
1127+
merged_events.append(
10761128
InvocationEvent(
10771129
author=e.author,
10781130
# The final response's text is already on
@@ -1085,7 +1137,7 @@ def convert_events_to_eval_invocations(
10851137
else None
10861138
),
10871139
grounding_metadata=e.grounding_metadata,
1088-
usage_metadata=e.usage_metadata,
1140+
usage_metadata=merged_usage.get(index, e.usage_metadata),
10891141
model_version=e.model_version,
10901142
)
10911143
)
@@ -1095,7 +1147,7 @@ def convert_events_to_eval_invocations(
10951147
user_content=user_content,
10961148
final_response=final_response,
10971149
intermediate_data=InvocationEvents(
1098-
invocation_events=invocation_events
1150+
invocation_events=merged_events
10991151
),
11001152
creation_timestamp=invocation_timestamp,
11011153
duration=(durations_per_invocation or {}).get(invocation_id),

0 commit comments

Comments
 (0)