fix: aclose all async generators to fix OTel tracing context

See https://github.com/google/adk-python/issues/1670#issuecomment-3115891100

PiperOrigin-RevId: 794659547
This commit is contained in:
Kacper Jawoszek
2025-08-13 11:18:26 -07:00
committed by Copybara-Service
parent c5af44cfc0
commit a30c63c593
23 changed files with 735 additions and 514 deletions
@@ -30,6 +30,7 @@ from ..runners import Runner
from ..sessions.base_session_service import BaseSessionService
from ..sessions.in_memory_session_service import InMemorySessionService
from ..sessions.session import Session
from ..utils.context_utils import Aclosing
from .eval_case import EvalCase
from .eval_case import IntermediateData
from .eval_case import Invocation
@@ -189,18 +190,25 @@ class EvaluationGenerator:
tool_uses = []
invocation_id = ""
async for event in runner.run_async(
user_id=user_id, session_id=session_id, new_message=user_content
):
invocation_id = (
event.invocation_id if not invocation_id else invocation_id
)
async with Aclosing(
runner.run_async(
user_id=user_id, session_id=session_id, new_message=user_content
)
) as agen:
async for event in agen:
invocation_id = (
event.invocation_id if not invocation_id else invocation_id
)
if event.is_final_response() and event.content and event.content.parts:
final_response = event.content
elif event.get_function_calls():
for call in event.get_function_calls():
tool_uses.append(call)
if (
event.is_final_response()
and event.content
and event.content.parts
):
final_response = event.content
elif event.get_function_calls():
for call in event.get_function_calls():
tool_uses.append(call)
response_invocations.append(
Invocation(