Skip to content

Commit 1a87114

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
refactor: Simplifying Tracing code
Replacing `getValidCurrentSpan.ifPresent` with a simpler `traceWithSpan` method. PiperOrigin-RevId: 881960190
1 parent e0d833b commit 1a87114

1 file changed

Lines changed: 99 additions & 103 deletions

File tree

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

Lines changed: 99 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -127,13 +127,13 @@ public class Tracing {
127127

128128
private Tracing() {}
129129

130-
private static Optional<Span> getValidCurrentSpan(String methodName) {
130+
private static void traceWithSpan(String methodName, Consumer<Span> traceAction) {
131131
Span span = Span.current();
132132
if (!span.getSpanContext().isValid()) {
133133
log.trace("{}: No valid span in current context.", methodName);
134-
return Optional.empty();
134+
return;
135135
}
136-
return Optional.of(span);
136+
traceAction.accept(span);
137137
}
138138

139139
private static void setInvocationAttributes(
@@ -206,16 +206,16 @@ public static void traceAgentInvocation(
206206
*/
207207
public static void traceToolCall(
208208
String toolName, String toolDescription, String toolType, Map<String, Object> args) {
209-
getValidCurrentSpan("traceToolCall")
210-
.ifPresent(
211-
span -> {
212-
setToolExecutionAttributes(span);
213-
span.setAttribute(GEN_AI_TOOL_NAME, toolName);
214-
span.setAttribute(GEN_AI_TOOL_DESCRIPTION, toolDescription);
215-
span.setAttribute(GEN_AI_TOOL_TYPE, toolType);
216-
217-
setJsonAttribute(span, ADK_TOOL_CALL_ARGS, args);
218-
});
209+
traceWithSpan(
210+
"traceToolCall",
211+
span -> {
212+
setToolExecutionAttributes(span);
213+
span.setAttribute(GEN_AI_TOOL_NAME, toolName);
214+
span.setAttribute(GEN_AI_TOOL_DESCRIPTION, toolDescription);
215+
span.setAttribute(GEN_AI_TOOL_TYPE, toolType);
216+
217+
setJsonAttribute(span, ADK_TOOL_CALL_ARGS, args);
218+
});
219219
}
220220

221221
/**
@@ -225,33 +225,33 @@ public static void traceToolCall(
225225
* @param functionResponseEvent The function response event.
226226
*/
227227
public static void traceToolResponse(String eventId, Event functionResponseEvent) {
228-
getValidCurrentSpan("traceToolResponse")
229-
.ifPresent(
230-
span -> {
231-
setToolExecutionAttributes(span);
232-
span.setAttribute(ADK_EVENT_ID, eventId);
233-
234-
FunctionResponse functionResponse =
235-
functionResponseEvent.functionResponses().stream().findFirst().orElse(null);
236-
237-
String toolCallId = "<not specified>";
238-
Object toolResponse = "<not specified>";
239-
if (functionResponse != null) {
240-
toolCallId = functionResponse.id().orElse(toolCallId);
241-
if (functionResponse.response().isPresent()) {
242-
toolResponse = functionResponse.response().get();
243-
}
244-
}
245-
246-
span.setAttribute(GEN_AI_TOOL_CALL_ID, toolCallId);
247-
248-
Object finalToolResponse =
249-
(toolResponse instanceof Map)
250-
? toolResponse
251-
: ImmutableMap.of("result", toolResponse);
252-
253-
setJsonAttribute(span, ADK_TOOL_RESPONSE, finalToolResponse);
254-
});
228+
traceWithSpan(
229+
"traceToolResponse",
230+
span -> {
231+
setToolExecutionAttributes(span);
232+
span.setAttribute(ADK_EVENT_ID, eventId);
233+
234+
FunctionResponse functionResponse =
235+
functionResponseEvent.functionResponses().stream().findFirst().orElse(null);
236+
237+
String toolCallId = "<not specified>";
238+
Object toolResponse = "<not specified>";
239+
if (functionResponse != null) {
240+
toolCallId = functionResponse.id().orElse(toolCallId);
241+
if (functionResponse.response().isPresent()) {
242+
toolResponse = functionResponse.response().get();
243+
}
244+
}
245+
246+
span.setAttribute(GEN_AI_TOOL_CALL_ID, toolCallId);
247+
248+
Object finalToolResponse =
249+
(toolResponse instanceof Map)
250+
? toolResponse
251+
: ImmutableMap.of("result", toolResponse);
252+
253+
setJsonAttribute(span, ADK_TOOL_RESPONSE, finalToolResponse);
254+
});
255255
}
256256

257257
/**
@@ -296,58 +296,54 @@ public static void traceCallLlm(
296296
String eventId,
297297
LlmRequest llmRequest,
298298
LlmResponse llmResponse) {
299-
getValidCurrentSpan("traceCallLlm")
300-
.ifPresent(
301-
span -> {
302-
span.setAttribute(GEN_AI_SYSTEM, "gcp.vertex.agent");
303-
llmRequest
304-
.model()
305-
.ifPresent(modelName -> span.setAttribute(GEN_AI_REQUEST_MODEL, modelName));
306-
307-
setInvocationAttributes(span, invocationContext, eventId);
308-
309-
setJsonAttribute(span, ADK_LLM_REQUEST, buildLlmRequestForTrace(llmRequest));
310-
setJsonAttribute(span, ADK_LLM_RESPONSE, llmResponse);
311-
312-
llmRequest
313-
.config()
314-
.ifPresent(
315-
config -> {
316-
config
317-
.topP()
318-
.ifPresent(
319-
topP ->
320-
span.setAttribute(GEN_AI_REQUEST_TOP_P, topP.doubleValue()));
321-
config
322-
.maxOutputTokens()
323-
.ifPresent(
324-
maxTokens ->
325-
span.setAttribute(
326-
GEN_AI_REQUEST_MAX_TOKENS, maxTokens.longValue()));
327-
});
328-
llmResponse
329-
.usageMetadata()
330-
.ifPresent(
331-
usage -> {
332-
usage
333-
.promptTokenCount()
334-
.ifPresent(
335-
tokens ->
336-
span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, (long) tokens));
337-
usage
338-
.candidatesTokenCount()
339-
.ifPresent(
340-
tokens ->
341-
span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, (long) tokens));
342-
});
343-
llmResponse
344-
.finishReason()
345-
.map(reason -> reason.knownEnum().name().toLowerCase(Locale.ROOT))
346-
.ifPresent(
347-
reason ->
348-
span.setAttribute(
349-
GEN_AI_RESPONSE_FINISH_REASONS, ImmutableList.of(reason)));
350-
});
299+
traceWithSpan(
300+
"traceCallLlm",
301+
span -> {
302+
span.setAttribute(GEN_AI_SYSTEM, "gcp.vertex.agent");
303+
llmRequest
304+
.model()
305+
.ifPresent(modelName -> span.setAttribute(GEN_AI_REQUEST_MODEL, modelName));
306+
307+
setInvocationAttributes(span, invocationContext, eventId);
308+
309+
setJsonAttribute(span, ADK_LLM_REQUEST, buildLlmRequestForTrace(llmRequest));
310+
setJsonAttribute(span, ADK_LLM_RESPONSE, llmResponse);
311+
312+
llmRequest
313+
.config()
314+
.ifPresent(
315+
config -> {
316+
config
317+
.topP()
318+
.ifPresent(
319+
topP -> span.setAttribute(GEN_AI_REQUEST_TOP_P, topP.doubleValue()));
320+
config
321+
.maxOutputTokens()
322+
.ifPresent(
323+
maxTokens ->
324+
span.setAttribute(
325+
GEN_AI_REQUEST_MAX_TOKENS, maxTokens.longValue()));
326+
});
327+
llmResponse
328+
.usageMetadata()
329+
.ifPresent(
330+
usage -> {
331+
usage
332+
.promptTokenCount()
333+
.ifPresent(
334+
tokens -> span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, (long) tokens));
335+
usage
336+
.candidatesTokenCount()
337+
.ifPresent(
338+
tokens -> span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, (long) tokens));
339+
});
340+
llmResponse
341+
.finishReason()
342+
.map(reason -> reason.knownEnum().name().toLowerCase(Locale.ROOT))
343+
.ifPresent(
344+
reason ->
345+
span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, ImmutableList.of(reason)));
346+
});
351347
}
352348

353349
/**
@@ -359,17 +355,17 @@ public static void traceCallLlm(
359355
*/
360356
public static void traceSendData(
361357
InvocationContext invocationContext, String eventId, List<Content> data) {
362-
getValidCurrentSpan("traceSendData")
363-
.ifPresent(
364-
span -> {
365-
setInvocationAttributes(span, invocationContext, eventId);
366-
367-
ImmutableList<Content> safeData =
368-
Optional.ofNullable(data).orElse(ImmutableList.of()).stream()
369-
.filter(Objects::nonNull)
370-
.collect(toImmutableList());
371-
setJsonAttribute(span, ADK_DATA, safeData);
372-
});
358+
traceWithSpan(
359+
"traceSendData",
360+
span -> {
361+
setInvocationAttributes(span, invocationContext, eventId);
362+
363+
ImmutableList<Content> safeData =
364+
Optional.ofNullable(data).orElse(ImmutableList.of()).stream()
365+
.filter(Objects::nonNull)
366+
.collect(toImmutableList());
367+
setJsonAttribute(span, ADK_DATA, safeData);
368+
});
373369
}
374370

375371
/**

0 commit comments

Comments
 (0)