Skip to content

fix: respect OpenTelemetry LIFO scoping in reactive streams #1516

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
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: respect OpenTelemetry LIFO scoping in reactive streams #1516

Show file tree
Hide file tree
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
49 changes: 29 additions & 20 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.google.adk.plugins.Plugin;
import com.google.adk.telemetry.Instrumentation;
import com.google.adk.telemetry.Instrumentation.AgentInvocation;
import com.google.adk.telemetry.Tracing;
import com.google.adk.utils.AgentEnums.AgentOrigin;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
Expand Down Expand Up @@ -331,31 +332,39 @@ private Flowable<Event> run(
},
agentInvocation -> {
InvocationContext invocationContext = agentInvocation.getCtx();
Context otelContext = agentInvocation.context().otelContext();
Flowable<Event> mainAndAfterEvents =
Flowable.defer(() -> runImplementation.apply(invocationContext))
.compose(Tracing.withContext(otelContext))
.concatWith(
Flowable.defer(
() ->
callCallback(
afterCallbacksToFunctions(
invocationContext.pluginManager(), afterAgentCallback),
invocationContext)
.toFlowable()));

return callCallback(
beforeCallbacksToFunctions(
invocationContext.pluginManager(), beforeAgentCallback),
invocationContext)
.flatMapPublisher(
beforeEvent -> {
if (invocationContext.endInvocation()) {
return Flowable.just(beforeEvent);
}
return Flowable.just(beforeEvent).concatWith(mainAndAfterEvents);
})
.switchIfEmpty(mainAndAfterEvents)
() ->
callCallback(
afterCallbacksToFunctions(
invocationContext.pluginManager(),
afterAgentCallback),
invocationContext)
.toFlowable())
.compose(Tracing.withContext(otelContext)));

return Flowable.defer(
() ->
callCallback(
beforeCallbacksToFunctions(
invocationContext.pluginManager(), beforeAgentCallback),
invocationContext)
.compose(Tracing.withContext(otelContext))
.flatMapPublisher(
beforeEvent -> {
if (invocationContext.endInvocation()) {
return Flowable.just(beforeEvent);
}
return Flowable.just(beforeEvent).concatWith(mainAndAfterEvents);
})
.switchIfEmpty(mainAndAfterEvents))
.doOnNext(agentInvocation::addEvent)
.doOnError(agentInvocation::setError);
.doOnError(agentInvocation::setError)
.compose(Tracing.withContext(otelContext));
},
AgentInvocation::close);
}
Expand Down
76 changes: 50 additions & 26 deletions core/src/main/java/com/google/adk/flows/llmflows/Functions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,25 +304,10 @@ private static Function<FunctionCall, Maybe<Event>> getFunctionCallMapper(
Map<String, Object> functionArgs =
functionCall.args().map(HashMap::new).orElse(new HashMap<>());

Maybe<Map<String, Object>> maybeFunctionResult =
maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
.switchIfEmpty(
Maybe.defer(
() ->
isLive
? processFunctionLive(
invocationContext,
tool,
toolContext,
functionCall,
functionArgs)
: callTool(tool, functionArgs, toolContext))
.compose(Tracing.withContext(parentContext)));

return postProcessFunctionResult(
maybeFunctionResult,
invocationContext,
tool,
functionCall,
functionArgs,
toolContext,
isLive,
Expand Down Expand Up @@ -487,9 +472,9 @@ static boolean hasPendingLongRunningCall(List<Event> events) {
}

private static Maybe<Event> postProcessFunctionResult(
Maybe<Map<String, Object>> maybeFunctionResult,
InvocationContext invocationContext,
BaseTool tool,
FunctionCall functionCall,
Map<String, Object> functionArgs,
ToolContext toolContext,
boolean isLive,
Expand All @@ -498,11 +483,38 @@ private static Maybe<Event> postProcessFunctionResult(
() ->
Instrumentation.recordToolExecution(
tool, invocationContext.agent(), functionArgs, parentContext),
toolExecution ->
processFunctionResult(
maybeFunctionResult, invocationContext, tool, functionArgs, toolContext, isLive)
.doOnSuccess(event -> toolExecution.context().setFunctionResponseEvent(event))
.doOnError(toolExecution::setError),
toolExecution -> {
Context toolOtelContext = toolExecution.context().otelContext();
Maybe<Map<String, Object>> maybeFunctionResult =
Maybe.defer(
() ->
maybeInvokeBeforeToolCall(
invocationContext, tool, functionArgs, toolContext))
.compose(Tracing.withContext(toolOtelContext))
.switchIfEmpty(
Maybe.defer(
() ->
isLive
? processFunctionLive(
invocationContext,
tool,
toolContext,
functionCall,
functionArgs)
: callTool(tool, functionArgs, toolContext))
.compose(Tracing.withContext(toolOtelContext)));
return processFunctionResult(
maybeFunctionResult,
invocationContext,
tool,
functionArgs,
toolContext,
isLive,
toolOtelContext)
.compose(Tracing.withContext(toolOtelContext))
.doOnSuccess(event -> toolExecution.context().setFunctionResponseEvent(event))
.doOnError(toolExecution::setError);
},
ToolExecution::close);
}

Expand All @@ -512,14 +524,19 @@ private static Maybe<Event> processFunctionResult(
BaseTool tool,
Map<String, Object> functionArgs,
ToolContext toolContext,
boolean isLive) {
boolean isLive,
Context toolOtelContext) {
return maybeFunctionResult
.map(Optional::of)
.defaultIfEmpty(Optional.empty())
.onErrorResumeNext(
t -> {
Maybe<Map<String, Object>> errorCallbackResult =
handleOnToolErrorCallback(invocationContext, tool, functionArgs, toolContext, t);
Maybe.defer(
() ->
handleOnToolErrorCallback(
invocationContext, tool, functionArgs, toolContext, t))
.compose(Tracing.withContext(toolOtelContext));
Maybe<Optional<Map<String, Object>>> mappedResult;
if (isLive) {
// In live mode, handle null results from the error callback gracefully.
Expand All @@ -535,8 +552,15 @@ private static Maybe<Event> processFunctionResult(
optionalInitialResult -> {
Map<String, Object> initialFunctionResult = optionalInitialResult.orElse(null);

return maybeInvokeAfterToolCall(
invocationContext, tool, functionArgs, toolContext, initialFunctionResult)
return Maybe.defer(
() ->
maybeInvokeAfterToolCall(
invocationContext,
tool,
functionArgs,
toolContext,
initialFunctionResult))
.compose(Tracing.withContext(toolOtelContext))
.map(Optional::of)
.defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
.flatMapMaybe(
Expand Down
24 changes: 10 additions & 14 deletions core/src/main/java/com/google/adk/telemetry/Instrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
Expand Down Expand Up @@ -91,9 +90,6 @@ public abstract static class ClosableTelemetryScope implements AutoCloseable {
/** The OpenTelemetry span associated with this scope. */
protected final Span span;

/** The OpenTelemetry scope associated with this span. */
protected final Scope scope;

/** The telemetry context for this scope. */
protected final TelemetryContext telemetryContext;

Expand All @@ -104,16 +100,15 @@ public abstract static class ClosableTelemetryScope implements AutoCloseable {
protected final AtomicBoolean closed = new AtomicBoolean(false);

/**
* Constructs a new {@code ClosableTelemetryScope} with the given span.
* Constructs a new {@code ClosableTelemetryScope} with the given span and parent context.
*
* @param span The OpenTelemetry span to manage.
* @param parentContext The OpenTelemetry parent context.
*/
@SuppressWarnings("MustBeClosedChecker")
ClosableTelemetryScope(Span span) {
ClosableTelemetryScope(Span span, Context parentContext) {
this.startTimeNanos = System.nanoTime();
this.span = span;
this.scope = span.makeCurrent();
this.telemetryContext = new TelemetryContext(Context.current());
this.telemetryContext = new TelemetryContext(parentContext.with(span));
}

/**
Expand All @@ -136,23 +131,22 @@ public void setError(Throwable caughtError) {
span.setStatus(StatusCode.ERROR, caughtError.getMessage());
}

/** Closes the scope and ends the underlying span, recording any applicable metrics. */
/** Ends the underlying span and records any applicable metrics. */
@Override
public final void close() {
if (closed.getAndSet(true)) {
return;
}
try {
beforeSpanEnd();
} finally {
span.end();
Duration elapsed = Duration.ofNanos(System.nanoTime() - startTimeNanos);
try {
recordMetrics(elapsed, caughtError);
} catch (RuntimeException e) {
handleMetricsError(e);
}
} finally {
scope.close();
}
}

Expand Down Expand Up @@ -184,7 +178,8 @@ public AgentInvocation(InvocationContext ctx, BaseAgent agent, Context parentCon
Tracing.getTracer()
.spanBuilder("invoke_agent " + agent.name())
.setParent(parentContext)
.startSpan());
.startSpan(),
parentContext);
this.agent = agent;
this.ctx = ctx;
Tracing.traceAgentInvocation(span, agent.name(), agent.description(), ctx);
Expand Down Expand Up @@ -254,7 +249,8 @@ public ToolExecution(
Tracing.getTracer()
.spanBuilder("execute_tool " + tool.name())
.setParent(parentContext)
.startSpan());
.startSpan(),
parentContext);
this.tool = tool;
this.agent = agent;
this.functionArgs = functionArgs;
Expand Down
Loading
Loading