Skip to content

Commit 20f863f

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
fix: Explicitly setting the otel parent spans in agents, llm flow and function calls
PiperOrigin-RevId: 881688036
1 parent 0d8e22d commit 20f863f

4 files changed

Lines changed: 32 additions & 28 deletions

File tree

‎core/src/main/java/com/google/adk/agents/BaseAgent.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3030
import com.google.errorprone.annotations.DoNotCall;
3131
import com.google.genai.types.Content;
32+
import io.opentelemetry.context.Context;
3233
import io.reactivex.rxjava3.core.Completable;
3334
import io.reactivex.rxjava3.core.Flowable;
3435
import io.reactivex.rxjava3.core.Maybe;
@@ -311,6 +312,7 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
311312
private Flowable<Event> run(
312313
InvocationContext parentContext,
313314
Function<InvocationContext, Flowable<Event>> runImplementation) {
315+
Context parentSpanContext = Context.current();
314316
return Flowable.defer(
315317
() -> {
316318
InvocationContext invocationContext = createInvocationContext(parentContext);
@@ -339,8 +341,12 @@ private Flowable<Event> run(
339341
})
340342
.switchIfEmpty(mainAndAfterEvents)
341343
.compose(
342-
Tracing.traceAgent(
343-
"invoke_agent " + name(), name(), description(), invocationContext));
344+
Tracing.<Event>trace("invoke_agent " + name())
345+
.setParent(parentSpanContext)
346+
.configure(
347+
span ->
348+
Tracing.traceAgentInvocation(
349+
span, name(), description(), invocationContext)));
344350
});
345351
}
346352

‎core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,10 @@ protected Flowable<Event> postprocess(
164164
* callbacks. Callbacks should not rely on its ID if they create their own separate events.
165165
*/
166166
private Flowable<LlmResponse> callLlm(
167-
InvocationContext context, LlmRequest llmRequest, Event eventForCallbackUsage) {
167+
Context spanContext,
168+
InvocationContext context,
169+
LlmRequest llmRequest,
170+
Event eventForCallbackUsage) {
168171
LlmAgent agent = (LlmAgent) context.agent();
169172

170173
LlmRequest.Builder llmRequestBuilder = llmRequest.toBuilder();
@@ -200,7 +203,7 @@ private Flowable<LlmResponse> callLlm(
200203
span.setStatus(StatusCode.ERROR, error.getMessage());
201204
span.recordException(error);
202205
})
203-
.compose(Tracing.trace("call_llm"))
206+
.compose(Tracing.<LlmResponse>trace("call_llm").setParent(spanContext))
204207
.concatMap(
205208
llmResp ->
206209
handleAfterModelCallback(context, llmResp, eventForCallbackUsage)
@@ -319,7 +322,7 @@ private Single<LlmResponse> handleAfterModelCallback(
319322
* @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
320323
* @throws IllegalStateException if a transfer agent is specified but not found.
321324
*/
322-
private Flowable<Event> runOneStep(InvocationContext context) {
325+
private Flowable<Event> runOneStep(Context spanContext, InvocationContext context) {
323326
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
324327

325328
return Flowable.defer(
@@ -351,7 +354,11 @@ private Flowable<Event> runOneStep(InvocationContext context) {
351354
.build();
352355
mutableEventTemplate.setTimestamp(0L);
353356

354-
return callLlm(context, llmRequestAfterPreprocess, mutableEventTemplate)
357+
return callLlm(
358+
spanContext,
359+
context,
360+
llmRequestAfterPreprocess,
361+
mutableEventTemplate)
355362
.concatMap(
356363
llmResponse -> {
357364
try (Scope postScope = currentContext.makeCurrent()) {
@@ -403,11 +410,12 @@ private Flowable<Event> runOneStep(InvocationContext context) {
403410
*/
404411
@Override
405412
public Flowable<Event> run(InvocationContext invocationContext) {
406-
return run(invocationContext, 0);
413+
return run(Context.current(), invocationContext, 0);
407414
}
408415

409-
private Flowable<Event> run(InvocationContext invocationContext, int stepsCompleted) {
410-
Flowable<Event> currentStepEvents = runOneStep(invocationContext).cache();
416+
private Flowable<Event> run(
417+
Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
418+
Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
411419
if (stepsCompleted + 1 >= maxSteps) {
412420
logger.debug("Ending flow execution because max steps reached.");
413421
return currentStepEvents;
@@ -427,7 +435,7 @@ private Flowable<Event> run(InvocationContext invocationContext, int stepsComple
427435
return Flowable.empty();
428436
} else {
429437
logger.debug("Continuing to next step of the flow.");
430-
return run(invocationContext, stepsCompleted + 1);
438+
return run(spanContext, invocationContext, stepsCompleted + 1);
431439
}
432440
}));
433441
}
@@ -444,6 +452,7 @@ private Flowable<Event> run(InvocationContext invocationContext, int stepsComple
444452
public Flowable<Event> runLive(InvocationContext invocationContext) {
445453
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
446454
Flowable<Event> preprocessEvents = preprocess(invocationContext, llmRequestRef);
455+
Context spanContext = Context.current();
447456

448457
return preprocessEvents.concatWith(
449458
Flowable.defer(
@@ -481,7 +490,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
481490
eventIdForSendData,
482491
llmRequestAfterPreprocess.contents());
483492
})
484-
.compose(Tracing.trace("send_data"));
493+
.compose(Tracing.<Event>trace("send_data").setParent(spanContext));
485494

486495
Flowable<LiveRequest> liveRequests =
487496
invocationContext

‎core/src/main/java/com/google/adk/flows/llmflows/Functions.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ public static Maybe<Event> handleFunctionCalls(
178178
if (events.size() > 1) {
179179
return Maybe.just(mergedEvent)
180180
.doOnSuccess(event -> Tracing.traceToolResponse(event.id(), event))
181-
.compose(Tracing.trace("tool_response", parentContext));
181+
.compose(Tracing.<Event>trace("tool_response").setParent(parentContext));
182182
}
183183
return Maybe.just(mergedEvent);
184184
});
@@ -432,8 +432,8 @@ private static Maybe<Event> postProcessFunctionResult(
432432
toolContext,
433433
invocationContext))
434434
.compose(
435-
Tracing.trace(
436-
"tool_response [" + tool.name() + "]", parentContext))
435+
Tracing.<Event>trace("tool_response [" + tool.name() + "]")
436+
.setParent(parentContext))
437437
.doOnSuccess(event -> Tracing.traceToolResponse(event.id(), event));
438438
});
439439
}
@@ -593,7 +593,9 @@ private static Maybe<Map<String, Object>> callTool(
593593
Tracing.traceToolCall(
594594
tool.name(), tool.description(), tool.getClass().getSimpleName(), args))
595595
.doOnError(t -> Span.current().recordException(t))
596-
.compose(Tracing.trace("tool_call [" + tool.name() + "]", parentContext))
596+
.compose(
597+
Tracing.<Map<String, Object>>trace("tool_call [" + tool.name() + "]")
598+
.setParent(parentContext))
597599
.onErrorResumeNext(
598600
e ->
599601
Maybe.error(

‎core/src/main/java/com/google/adk/telemetry/Tracing.java‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -426,19 +426,6 @@ public static <T> TracerProvider<T> trace(String spanName) {
426426
return new TracerProvider<>(spanName);
427427
}
428428

429-
/**
430-
* Returns a transformer that traces the execution of an RxJava stream with an explicit parent
431-
* context.
432-
*
433-
* @param spanName The name of the span to create.
434-
* @param parentContext The explicit parent context for the span.
435-
* @param <T> The type of the stream.
436-
* @return A TracerProvider that can be used with .compose().
437-
*/
438-
public static <T> TracerProvider<T> trace(String spanName, Context parentContext) {
439-
return new TracerProvider<T>(spanName).setParent(parentContext);
440-
}
441-
442429
/**
443430
* Returns a transformer that traces an agent invocation.
444431
*

0 commit comments

Comments
 (0)