Skip to content

fix: ensure call_llm tracing spans are always ended in multi-agent setups #4742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
atian8179 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: ensure call_llm tracing spans are always ended in multi-agent setups #4742

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion src/google/adk/flows/llm_flows/base_llm_flow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@
from ...models.llm_request import LlmRequest
from ...models.llm_response import LlmResponse
from ...telemetry import tracing
from opentelemetry import context as otel_context
from opentelemetry import trace

from ...telemetry.tracing import trace_call_llm
from ...telemetry.tracing import trace_send_data
from ...telemetry.tracing import tracer
Expand Down Expand Up @@ -1127,7 +1130,14 @@ async def _call_llm_async(
llm = self.__get_llm(invocation_context)

async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
with tracer.start_as_current_span('call_llm') as span:
# Use explicit span management instead of start_as_current_span context
# manager to ensure span.end() is always called even when GeneratorExit
# is raised during async iteration (e.g., when transfer_to_agent causes
# the generator to be closed). See #4715.
span = tracer.start_span('call_llm')
ctx = trace.set_span_in_context(span)
token = otel_context.attach(ctx)
try:
if invocation_context.run_config.support_cfc:
invocation_context.live_request_queue = LiveRequestQueue()
responses_generator = self.run_live(invocation_context)
Expand Down Expand Up @@ -1187,6 +1197,12 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
llm_response = altered_llm_response

yield llm_response
finally:
try:
otel_context.detach(token)
except ValueError:
pass
span.end()

async with Aclosing(_call_llm_with_tracing()) as agen:
async for event in agen:
Expand Down