Skip to content

feat(profiling): register llm.agent.phase slot and add phase set/clear API #11529

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

Draft
jbachorik wants to merge 1 commit into
base: jb/llm-ddt-1-config
Choose a base branch
from
Draft

feat(profiling): register llm.agent.phase slot and add phase set/clear API #11529

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
15 changes: 15 additions & 0 deletions ...rofiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isCpuProfilerEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isLiveHeapSizeTrackingEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isMemoryLeakProfilingEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isLlmPhaseAttributeEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isResourceNameContextAttributeEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isSpanNameContextAttributeEnabled;
import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isTrackingGenerations;
Expand Down Expand Up @@ -69,6 +70,7 @@ public final class DatadogProfiler {

private static final String OPERATION = "_dd.trace.operation";
private static final String RESOURCE = "_dd.trace.resource";
private static final String LLM_PHASE = "llm.agent.phase";

private static final int MAX_NUM_ENDPOINTS = 8192;

Expand Down Expand Up @@ -105,6 +107,7 @@ public static DatadogProfiler newInstance(ConfigProvider configProvider) {
private final Set<ProfilingMode> profilingModes = EnumSet.noneOf(ProfilingMode.class);

private final ContextSetter contextSetter;
private final int llmPhaseOffset;

private final List<String> orderedContextAttributes;

Expand Down Expand Up @@ -150,7 +153,11 @@ private DatadogProfiler(ConfigProvider configProvider) {
if (isResourceNameContextAttributeEnabled(configProvider)) {
orderedContextAttributes.add(RESOURCE);
}
if (isLlmPhaseAttributeEnabled(configProvider)) {
orderedContextAttributes.add(LLM_PHASE);
}
this.contextSetter = new ContextSetter(profiler, orderedContextAttributes);
this.llmPhaseOffset = contextSetter.offsetOf(LLM_PHASE);
this.queueTimeThresholdMillis =
configProvider.getLong(
PROFILING_QUEUEING_TIME_THRESHOLD_MILLIS,
Expand Down Expand Up @@ -408,6 +415,14 @@ public boolean clearContextValue(int offset) {
return false;
}

public boolean setAgentPhase(String phaseToken) {
return contextSetter.setContextValue(llmPhaseOffset, phaseToken);
}

public boolean clearAgentPhase() {
return contextSetter.clearContextValue(llmPhaseOffset);
}

private void debugLogging(long localRootSpanId) {
if (detailedDebugLogging && log.isDebugEnabled()) {
log.debug("localRootSpanId={}", localRootSpanId, new Throwable());
Expand Down
Loading