diff --git a/data-prepper-plugins/otel-trace-raw-processor/README.md b/data-prepper-plugins/otel-trace-raw-processor/README.md index b64b074621..de3ee03b68 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/README.md +++ b/data-prepper-plugins/otel-trace-raw-processor/README.md @@ -23,8 +23,16 @@ This processor automatically enriches GenAI agent traces. The enrichment is alwa **What it does:** 1. **Vendor attribute normalization** — Maps vendor-specific attributes (OpenInference, OpenLLMetry) to [OTel GenAI Semantic Conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/). Original attributes are preserved. See [`genai-attribute-mappings.yaml`](src/main/resources/genai-attribute-mappings.yaml) for the full mapping table. -2. **Root span enrichment** — Propagates `gen_ai.system`, `gen_ai.provider.name`, `gen_ai.agent.name`, `gen_ai.request.model`, and `gen_ai.operation.name` from child spans to root (first-child-wins, skip-if-present). Aggregates `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` across children (sum). -3. **Flattened sub-key stripping** — Removes flattened sub-keys (e.g. `llm.input_messages.0.message.content`) that conflict with parent string values, preventing OpenSearch mapping failures. +2. **Root span enrichment** — Propagates `gen_ai.system`, `gen_ai.provider.name`, `gen_ai.agent.name`, and `gen_ai.request.model` from child spans to root (skip-if-present). Values come from the outermost GenAI span first (earliest start time), then other GenAI spans, then the remaining children. Aggregates `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` across children, counting only spans that have no ancestor reporting usage: agent spans such as `invoke_agent` often report the total of their own LLM calls, which would otherwise be counted twice. This assumes an ancestor's reported usage includes its subtree; a sub-agent that runs its own LLM calls under a parent reporting only its own usage is not added. +3. **Agent root marker** — Sets `derived.gen_ai.is_root: true` on each outermost GenAI span of a trace: a span with `gen_ai.operation.name` and no GenAI ancestor. That is the root itself when it is a GenAI span, or the agent span below a non-GenAI root (for example an agent served behind an HTTP handler). It also aggregates the token usage of that span's descendants onto it (with the same rule), unless it already reports `gen_ai.usage.input_tokens` or `gen_ai.usage.output_tokens`. + * Several top-level agents under one root are each marked. + * An agent that calls another agent over HTTP or RPC in the same trace (`invoke_agent` > HTTP client > HTTP server > `invoke_agent`) gets one mark, on the calling agent, because the called agent has a GenAI ancestor. If the called service's spans are flushed after `trace_flush_interval` without the caller's spans, the called agent is marked too, and consumers that list marked spans show it as a separate entry. + * Spans that arrive after their trace's root has been processed are enriched on their own. The processor remembers, for twice `trace_flush_interval` plus `trace_group_cache_ttl`, which span ids of a trace are already covered by a GenAI span or a marker, so such a span is not mistaken for the outermost one, including spans flushed without their root. A span that arrives later than that, or after the cache (500,000 span ids, see the `genAiCoverageCount` metric) has evicted its trace, may be marked as well. + * If a trace's children are flushed before a GenAI root arrives, both the root and the outermost child can end up marked. Consumers should prefer the marked span with no parent. + * A marker set at the source is kept. A span set to `false` is never marked; set it on a trace's outermost GenAI spans to opt the trace out. + * Spans flushed after `trace_flush_interval` without their root are still marked, because the marker only needs the spans below it. + * If any span in the trace already carries `derived.gen_ai.is_root: true` (set by an SDK or by hand), no spans are marked by the processor. To choose the root yourself, set it to `true` on that span when you create it. +4. **Flattened sub-key stripping** — Removes flattened sub-keys (e.g. `llm.input_messages.0.message.content`) that conflict with parent string values, preventing OpenSearch mapping failures. **Important:** The enrichment matches attributes by their original OTel key names (e.g. `http.method`). This requires the source to use `output_format: otel` so that attribute keys are not transformed. The default `opensearch` format rewrites keys (e.g. `http.method` → `span.attributes.http@method`), which prevents the enrichment from working. See the [otel-trace-source README](../otel-trace-source/README.md#output-format) for details. diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverage.java b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverage.java new file mode 100644 index 0000000000..d11fea107d --- /dev/null +++ b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverage.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.oteltrace; + +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; + +import java.time.Duration; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import java.util.function.Function; + +/** + * Remembers, per trace, the span ids already covered by a GenAI span or an agent root marker (the + * span itself or one of its ancestors). Spans that arrive after their trace's root has been + * processed are enriched without their ancestors; this lets the agent root marker tell whether + * such a span is really outermost. + */ +final class GenAiCoverage { + + /** Keeps nothing: every parent outside the batch counts as uncovered. */ + static final GenAiCoverage NONE = new GenAiCoverage(null); + + private final Cache> coveredByTrace; + + private GenAiCoverage(final Cache> coveredByTrace) { + this.coveredByTrace = coveredByTrace; + } + + static GenAiCoverage create(final long maxSpanIds, final Duration timeToLive) { + return new GenAiCoverage(Caffeine.newBuilder() + .maximumWeight(maxSpanIds) + .>weigher((traceId, spanIds) -> spanIds.size()) + .expireAfterWrite(timeToLive) + .build()); + } + + /** + * Runs {@code enrich} with the trace's span ids covered so far and records the ids it returns. + * The read, the enrichment and the write happen atomically for the trace, so pipeline workers + * processing spans of one trace at the same time cannot both take a span for the outermost. + */ + void update(final String traceId, final Function, Set> enrich) { + if (coveredByTrace == null || traceId == null) { + enrich.apply(Collections.emptySet()); + return; + } + coveredByTrace.asMap().compute(traceId, (id, existing) -> { + final Set known = existing == null ? Collections.emptySet() : existing; + final Set added = enrich.apply(known); + if (added.isEmpty() || known.containsAll(added)) { + return existing; + } + final Set merged = new HashSet<>(known); + merged.addAll(added); + return merged; + }); + } + + /** Span ids of the trace covered so far; for tests. */ + Set covered(final String traceId) { + if (coveredByTrace == null || traceId == null) { + return Collections.emptySet(); + } + final Set spanIds = coveredByTrace.getIfPresent(traceId); + return spanIds == null ? Collections.emptySet() : spanIds; + } + + long estimatedSize() { + return coveredByTrace == null ? 0 : coveredByTrace.estimatedSize(); + } +} diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiEnrichmentHelper.java b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiEnrichmentHelper.java index 301df17d5f..d4bd510271 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiEnrichmentHelper.java +++ b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiEnrichmentHelper.java @@ -14,11 +14,18 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; +import java.time.format.DateTimeParseException; import java.util.ArrayList; import java.util.Collection; +import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.function.Predicate; /** * Enriches GenAI agent traces by propagating attributes from child spans to root @@ -35,6 +42,14 @@ final class GenAiEnrichmentHelper { private static final String GEN_AI_REQUEST_MODEL_KEY = "gen_ai.request.model"; private static final String GEN_AI_INPUT_TOKENS_KEY = "gen_ai.usage.input_tokens"; private static final String GEN_AI_OUTPUT_TOKENS_KEY = "gen_ai.usage.output_tokens"; + private static final String GEN_AI_OPERATION_NAME_KEY = "gen_ai.operation.name"; + /** + * Marks the outermost GenAI span of a trace: the root when it is a GenAI span, else the agent + * span below a non-GenAI root (for example an agent behind an HTTP handler), so consumers can + * list the trace by that span. A source sets it to {@code true} to choose the AI roots itself, + * or to {@code false} to keep a span from being marked. + */ + static final String DERIVED_GEN_AI_IS_ROOT_KEY = "derived.gen_ai.is_root"; private static final String ATTRIBUTES_PREFIX = "attributes/"; private static final String[] PROPAGATED_STRING_KEYS = { @@ -50,100 +65,370 @@ private GenAiEnrichmentHelper() {} /** * Enriches a batch of spans: normalizes vendor attributes, strips flattened sub-keys, - * then propagates select gen_ai attributes from children to root spans grouped by traceId. + * marks the outermost GenAI spans of each trace, then + * propagates select gen_ai attributes from children to root spans grouped by traceId. */ static void enrichBatch(final List spans) { + enrichBatch(spans, GenAiCoverage.NONE); + } + + /** + * Same as {@link #enrichBatch(List)}, using and updating {@code coverage} so spans that arrive + * after their ancestors are not mistaken for a trace's outermost GenAI span. + */ + static void enrichBatch(final List spans, final GenAiCoverage coverage) { for (final Span span : spans) { normalizeAttributes(span); stripFlattenedSubkeys(span); } - final Map> spansByTrace = new HashMap<>(); + final Map> spansByTrace = new LinkedHashMap<>(); for (final Span span : spans) { spansByTrace.computeIfAbsent(span.getTraceId(), k -> new ArrayList<>()).add(span); } - for (final List traceSpans : spansByTrace.values()) { - Span rootSpan = null; - final List children = new ArrayList<>(); + for (final Map.Entry> trace : spansByTrace.entrySet()) { + final TraceIndex index = new TraceIndex(trace.getValue()); + if (!index.hasGenAiData()) { + continue; + } + // Read, mark and record under the trace's lock: spans of one trace can be in batches + // that pipeline workers process at the same time. + coverage.update(trace.getKey(), known -> enrichTrace(index, known)); + } + } + + /** Enriches one trace and returns the span ids it covers, for later batches of the trace. */ + private static Set enrichTrace(final TraceIndex index, final Set known) { + final List outermost = index.outermostGenAiSpans(known); + // Root first: it sums the spans' own tokens, before marking adds aggregates to agent roots. + final SpanInfo root = index.root(); + if (root != null && index.size() > 1) { + enrichRootSpan(index, root, known); + } + markAgentRoots(index, outermost); + return index.coveredSpanIds(known); + } + + /** One span with the attributes this helper reads, deserialized once. */ + private static final class SpanInfo { + private final Span span; + private final Map attributes; + private final boolean genAi; + private final Boolean marker; + private final Long inputTokens; + private final Long outputTokens; + private final Instant startTime; + private boolean marked; + + private SpanInfo(final Span span) { + this.span = span; + final Map attrs = span.getAttributes(); + this.attributes = attrs == null ? Map.of() : attrs; + this.genAi = attributes.get(GEN_AI_OPERATION_NAME_KEY) != null; + this.marker = toBoolean(attributes.get(DERIVED_GEN_AI_IS_ROOT_KEY)); + this.marked = Boolean.TRUE.equals(marker); + this.inputTokens = toLong(attributes.get(GEN_AI_INPUT_TOKENS_KEY)); + this.outputTokens = toLong(attributes.get(GEN_AI_OUTPUT_TOKENS_KEY)); + this.startTime = toInstant(span.getStartTime()); + } + + private String spanId() { + return span.getSpanId(); + } + + private String parentSpanId() { + return span.getParentSpanId(); + } + + private boolean hasTokens() { + return inputTokens != null || outputTokens != null; + } + + private boolean isRoot() { + final String parentSpanId = parentSpanId(); + return parentSpanId == null || parentSpanId.isEmpty() || "0000000000000000".equals(parentSpanId); + } + } + + private static final Comparator BY_START_TIME = Comparator.comparing( + (SpanInfo info) -> info.startTime, Comparator.nullsLast(Comparator.naturalOrder())); + + /** The spans of one trace in a batch, indexed once. */ + private static final class TraceIndex { + private final List spans = new ArrayList<>(); + private final Map bySpanId = new HashMap<>(); + private final Map> childrenByParent = new HashMap<>(); + + private TraceIndex(final Collection traceSpans) { for (final Span span : traceSpans) { - if (isRootSpan(span)) { - rootSpan = span; - } else { - children.add(span); + final SpanInfo info = new SpanInfo(span); + spans.add(info); + if (info.spanId() != null) { + bySpanId.put(info.spanId(), info); + } + if (info.parentSpanId() != null) { + childrenByParent.computeIfAbsent(info.parentSpanId(), k -> new ArrayList<>()).add(info); + } + } + } + + private int size() { + return spans.size(); + } + + private boolean hasGenAiData() { + for (final SpanInfo info : spans) { + if (info.genAi || info.marker != null || info.hasTokens()) { + return true; + } + for (final String key : PROPAGATED_STRING_KEYS) { + if (info.attributes.containsKey(key)) { + return true; + } + } + } + return false; + } + + private SpanInfo root() { + for (final SpanInfo info : spans) { + if (info.isRoot()) { + return info; } } - if (rootSpan != null && !children.isEmpty()) { - enrichRootSpan(rootSpan, children); + return null; + } + + private boolean hasSourceMarker() { + for (final SpanInfo info : spans) { + if (Boolean.TRUE.equals(info.marker)) { + return true; + } } + return false; + } + + /** + * Whether an ancestor of {@code info} matches: looked up in the batch, then, for the first + * parent outside it, in {@code known} (span ids covered in earlier batches). + */ + private boolean hasAncestor(final SpanInfo info, final Predicate matches, final Set known) { + final Set seen = new HashSet<>(); + String parentId = info.parentSpanId(); + SpanInfo parent = bySpanId.get(parentId); + while (parent != null && seen.add(parent.spanId())) { + if (matches.test(parent)) { + return true; + } + parentId = parent.parentSpanId(); + parent = bySpanId.get(parentId); + } + return parentId != null && known.contains(parentId); + } + + private boolean hasGenAiAncestor(final SpanInfo info, final Set known) { + return hasAncestor(info, parent -> parent.genAi || parent.marked, known); + } + + /** + * GenAI spans with no GenAI ancestor, earliest first. Ancestors are looked up in the batch, + * then in {@code known}; any other parent counts as a non-GenAI span. + */ + private List outermostGenAiSpans(final Set known) { + final List outermost = new ArrayList<>(); + for (final SpanInfo info : spans) { + if (info.genAi && !hasGenAiAncestor(info, known)) { + outermost.add(info); + } + } + outermost.sort(BY_START_TIME); + return outermost; + } + + /** Span ids of the batch that are, or are below, a GenAI span or a marked span. */ + private Set coveredSpanIds(final Set known) { + final Set covered = new HashSet<>(); + for (final SpanInfo info : spans) { + if (info.spanId() != null && (info.genAi || info.marked || hasGenAiAncestor(info, known))) { + covered.add(info.spanId()); + } + } + return covered; + } + + private List descendantsOf(final SpanInfo ancestor) { + final List descendants = new ArrayList<>(); + final List pending = new ArrayList<>(childrenByParent.getOrDefault(ancestor.spanId(), List.of())); + final Set seen = new HashSet<>(); + seen.add(ancestor.spanId()); + while (!pending.isEmpty()) { + final SpanInfo descendant = pending.remove(pending.size() - 1); + if (!seen.add(descendant.spanId())) { + continue; + } + descendants.add(descendant); + pending.addAll(childrenByParent.getOrDefault(descendant.spanId(), List.of())); + } + return descendants; + } + + /** + * Sums token usage over {@code subset}, counting only spans with no ancestor in the subset + * that also reports usage. Agent spans such as invoke_agent often report the total of their + * own LLM calls, so adding every span with usage would count those calls twice. + */ + private TokenTotals sumTopmostTokens(final Collection subset) { + final Set members = new HashSet<>(subset); + final TokenTotals totals = new TokenTotals(); + for (final SpanInfo info : subset) { + if (!info.hasTokens() + || hasAncestor(info, parent -> members.contains(parent) && parent.hasTokens(), Set.of())) { + continue; + } + totals.found = true; + totals.input += info.inputTokens != null ? info.inputTokens : 0; + totals.output += info.outputTokens != null ? info.outputTokens : 0; + } + return totals; } } - private static boolean isRootSpan(final Span span) { - final String parentSpanId = span.getParentSpanId(); - return parentSpanId == null || parentSpanId.isEmpty() - || "0000000000000000".equals(parentSpanId); + private static final class TokenTotals { + private long input; + private long output; + private boolean found; } /** - * Propagates gen_ai.* string attributes and aggregated token counts from children to root. - * Skips if root already has the attributes. + * Marks the outermost GenAI spans of a trace with {@code derived.gen_ai.is_root} and + * aggregates their descendants' tokens onto them. Skipped when a span is already marked + * {@code true}, so a span marked by the SDK or the user is the only AI root; a span set to + * {@code false} is never marked. */ - static void enrichRootSpan(final Span rootSpan, final Collection children) { - final Map rootAttrs = rootSpan.getAttributes(); + private static void markAgentRoots(final TraceIndex index, final List outermost) { + if (index.hasSourceMarker()) { + LOG.debug("Trace {} already has a span marked {} = true; not marking its spans", + index.spans.get(0).span.getTraceId(), DERIVED_GEN_AI_IS_ROOT_KEY); + return; + } + for (final SpanInfo agentRoot : outermost) { + if (Boolean.FALSE.equals(agentRoot.marker)) { + continue; + } + agentRoot.span.put(ATTRIBUTES_PREFIX + DERIVED_GEN_AI_IS_ROOT_KEY, true); + agentRoot.marked = true; + LOG.debug("Marked span {} as the GenAI root of trace {}", agentRoot.spanId(), agentRoot.span.getTraceId()); + if (agentRoot.hasTokens()) { + continue; + } + final TokenTotals tokens = index.sumTopmostTokens(index.descendantsOf(agentRoot)); + if (tokens.found) { + agentRoot.span.put(ATTRIBUTES_PREFIX + GEN_AI_INPUT_TOKENS_KEY, tokens.input); + agentRoot.span.put(ATTRIBUTES_PREFIX + GEN_AI_OUTPUT_TOKENS_KEY, tokens.output); + } + } + } - final Map toPropagate = new HashMap<>(); + /** + * Propagates gen_ai.* string attributes and aggregated token counts from children to root. + * Skips the attributes the root already has. Values come from the outermost GenAI spans + * first, so the trace gets the outer agent's name and model whatever order the batch is in. + */ + private static void enrichRootSpan(final TraceIndex index, final SpanInfo root, final Set known) { + final Map toPropagate = new LinkedHashMap<>(); for (final String key : PROPAGATED_STRING_KEYS) { - if (rootAttrs == null || !rootAttrs.containsKey(key)) { + if (!root.attributes.containsKey(key)) { toPropagate.put(key, null); } } - final boolean rootHasTokens = rootAttrs != null && rootAttrs.containsKey(GEN_AI_INPUT_TOKENS_KEY); - - if (toPropagate.isEmpty() && rootHasTokens) { + if (toPropagate.isEmpty() && root.hasTokens()) { return; } - long totalInputTokens = 0; - long totalOutputTokens = 0; - boolean foundTokens = false; - - for (final Span child : children) { - final Map attrs = child.getAttributes(); - if (attrs == null) { - continue; + final List children = new ArrayList<>(); + for (final SpanInfo info : index.spans) { + if (info != root) { + children.add(info); } - - for (final Map.Entry entry : toPropagate.entrySet()) { - if (entry.getValue() == null && attrs.containsKey(entry.getKey())) { - entry.setValue((String) attrs.get(entry.getKey())); + } + if (!toPropagate.isEmpty()) { + // Outermost among the children: GenAI spans with no GenAI ancestor below the root. + final Set outer = new HashSet<>(); + for (final SpanInfo child : children) { + if (child.genAi && !index.hasAncestor(child, parent -> parent != root && (parent.genAi || parent.marked), known)) { + outer.add(child); } } - - final Number inputTokens = (Number) attrs.get(GEN_AI_INPUT_TOKENS_KEY); - final Number outputTokens = (Number) attrs.get(GEN_AI_OUTPUT_TOKENS_KEY); - if (inputTokens != null || outputTokens != null) { - foundTokens = true; - if (inputTokens != null) { - totalInputTokens += inputTokens.longValue(); + final List ordered = new ArrayList<>(children); + ordered.sort(Comparator.comparingInt(info -> outer.contains(info) ? 0 : info.genAi ? 1 : 2) + .thenComparing(BY_START_TIME)); + for (final SpanInfo child : ordered) { + for (final Map.Entry entry : toPropagate.entrySet()) { + final Object value = child.attributes.get(entry.getKey()); + if (entry.getValue() == null && value != null) { + entry.setValue(value.toString()); + } } - if (outputTokens != null) { - totalOutputTokens += outputTokens.longValue(); + } + for (final Map.Entry entry : toPropagate.entrySet()) { + if (entry.getValue() != null) { + root.span.put(ATTRIBUTES_PREFIX + entry.getKey(), entry.getValue()); + LOG.debug("Propagated {} = {} to root span {}", entry.getKey(), entry.getValue(), root.spanId()); } } } - for (final Map.Entry entry : toPropagate.entrySet()) { - if (entry.getValue() != null) { - rootSpan.put(ATTRIBUTES_PREFIX + entry.getKey(), entry.getValue()); - LOG.debug("Propagated {} = {} to root span {}", entry.getKey(), entry.getValue(), rootSpan.getSpanId()); + if (!root.hasTokens()) { + final TokenTotals tokens = index.sumTopmostTokens(children); + if (tokens.found) { + root.span.put(ATTRIBUTES_PREFIX + GEN_AI_INPUT_TOKENS_KEY, tokens.input); + root.span.put(ATTRIBUTES_PREFIX + GEN_AI_OUTPUT_TOKENS_KEY, tokens.output); + LOG.debug("Aggregated tokens (input={}, output={}) to root span {}", tokens.input, tokens.output, root.spanId()); } } + } - if (!rootHasTokens && foundTokens) { - rootSpan.put(ATTRIBUTES_PREFIX + GEN_AI_INPUT_TOKENS_KEY, totalInputTokens); - rootSpan.put(ATTRIBUTES_PREFIX + GEN_AI_OUTPUT_TOKENS_KEY, totalOutputTokens); - LOG.debug("Aggregated tokens (input={}, output={}) to root span {}", totalInputTokens, totalOutputTokens, rootSpan.getSpanId()); + /** A token count from a number or a numeric string; anything else is treated as absent. */ + private static Long toLong(final Object value) { + if (value instanceof Number) { + return ((Number) value).longValue(); + } + if (value instanceof String) { + try { + return (long) Double.parseDouble(((String) value).trim()); + } catch (final NumberFormatException e) { + return null; + } + } + return null; + } + + private static Boolean toBoolean(final Object value) { + if (value instanceof Boolean) { + return (Boolean) value; + } + if (value instanceof String) { + final String text = ((String) value).trim(); + if ("true".equalsIgnoreCase(text)) { + return Boolean.TRUE; + } + if ("false".equalsIgnoreCase(text)) { + return Boolean.FALSE; + } + } + return null; + } + + /** Start times as instants: ISO-8601 strings with 0 to 9 fractional digits do not sort as text. */ + private static Instant toInstant(final String time) { + if (time == null) { + return null; + } + try { + return Instant.parse(time); + } catch (final DateTimeParseException e) { + return null; } } diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java index a8773a34da..50f8041f23 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java +++ b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java @@ -28,6 +28,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Duration; import java.util.Collection; import java.util.Collections; import java.util.Iterator; @@ -48,13 +49,18 @@ public class OTelTraceRawProcessor extends AbstractProcessor, Recor private static final Logger LOG = LoggerFactory.getLogger(OTelTraceRawProcessor.class); public static final String TRACE_GROUP_CACHE_COUNT_METRIC_NAME = "traceGroupCacheCount"; public static final String SPAN_SET_COUNT_METRIC_NAME = "spanSetCount"; + public static final String GENAI_COVERAGE_COUNT_METRIC_NAME = "genAiCoverageCount"; private final long traceFlushInterval; + static final long GENAI_COVERAGE_MAX_SPAN_IDS = 500_000L; private final Map traceIdSpanSetMap = new ConcurrentHashMap<>(); private final Cache traceIdTraceGroupCache; + /** Span ids covered by GenAI spans, per trace, for spans that arrive after their ancestors. */ + private final GenAiCoverage genAiCoverage; + private long lastTraceFlushTime = 0L; private final ReentrantLock traceFlushLock = new ReentrantLock(); @@ -74,8 +80,13 @@ public OTelTraceRawProcessor(final OtelTraceRawProcessorConfig otelTraceRawProce .build(); + // Outlives the orphan flush: spans that arrive late are held up to two flush intervals. + genAiCoverage = GenAiCoverage.create(GENAI_COVERAGE_MAX_SPAN_IDS, + Duration.ofMillis(2 * traceFlushInterval).plus(otelTraceRawProcessorConfig.getTraceGroupCacheTimeToLive())); + pluginMetrics.gauge(TRACE_GROUP_CACHE_COUNT_METRIC_NAME, traceIdTraceGroupCache, cache -> (double) cache.estimatedSize()); pluginMetrics.gauge(SPAN_SET_COUNT_METRIC_NAME, traceIdSpanSetMap, cache -> (double) cache.size()); + pluginMetrics.gauge(GENAI_COVERAGE_COUNT_METRIC_NAME, genAiCoverage, coverage -> (double) coverage.estimatedSize()); LOG.info("Configured Trace Raw Processor with a trace flush interval of {} ms.", traceFlushInterval); } @@ -100,7 +111,7 @@ public Collection> doExecute(Collection> records) { processedSpans.addAll(getTracesToFlushByGarbageCollection()); // Enrich GenAI agent traces (propagate select gen_ai attributes to root, aggregate tokens, strip conflicting sub-keys) - GenAiEnrichmentHelper.enrichBatch(processedSpans); + GenAiEnrichmentHelper.enrichBatch(processedSpans, genAiCoverage); // Derive server span attributes (fault, error, operation, environment) OTelSpanDerivationUtil.deriveServerSpanAttributes(processedSpans); diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiAgentRootTest.java b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiAgentRootTest.java new file mode 100644 index 0000000000..9f55c3dddd --- /dev/null +++ b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiAgentRootTest.java @@ -0,0 +1,448 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.oteltrace; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PipelineDescription; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; +import org.opensearch.dataprepper.model.trace.JacksonSpan; +import org.opensearch.dataprepper.model.trace.Span; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.oteltrace.GenAiEnrichmentHelper.DERIVED_GEN_AI_IS_ROOT_KEY; + +/** + * The outermost GenAI span of each trace is marked as its AI root: the root itself when it is a + * GenAI span, else the agent span below a non-GenAI root (an agent behind an HTTP handler). + */ +class GenAiAgentRootTest { + + private static final String TRACE_ID = "AGENT_ROOT_TRACE"; + + @Test + void marksTheOutermostGenAiSpanUnderAnHttpRoot() { + final List spans = httpAgentTrace(); + Collections.shuffle(spans); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agent"), markedSpanIds(spans)); + final Span http = byId(spans, "http"); + assertNull(http.getAttributes().get("gen_ai.operation.name")); + assertNull(http.getAttributes().get(DERIVED_GEN_AI_IS_ROOT_KEY)); + } + + @Test + void aggregatesDescendantTokensOntoTheAgentRootWithoutDoubleCountingTheTraceRoot() { + final List spans = httpAgentTrace(); + + GenAiEnrichmentHelper.enrichBatch(spans); + + final Map agent = byId(spans, "agent").getAttributes(); + assertEquals(15L, ((Number) agent.get("gen_ai.usage.input_tokens")).longValue()); + assertEquals(7L, ((Number) agent.get("gen_ai.usage.output_tokens")).longValue()); + final Map http = byId(spans, "http").getAttributes(); + assertEquals(15L, ((Number) http.get("gen_ai.usage.input_tokens")).longValue()); + assertEquals(7L, ((Number) http.get("gen_ai.usage.output_tokens")).longValue()); + } + + @Test + void propagatesTheOuterAgentNameToTheTraceRootWhateverTheBatchOrder() { + final List spans = httpAgentTrace(); + spans.add(span("inner", "chat1", "2026-01-01T00:00:03Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Weather Agent"))); + // Inner agent first in the batch: first-child-wins would pick it. + Collections.reverse(spans); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals("Travel Planner", byId(spans, "http").getAttributes().get("gen_ai.agent.name")); + } + + @Test + void countsAnAgentsReportedTotalOnceWhenItsCallsAlsoReportUsage() { + // Strands-style: invoke_agent reports the total of its own chat calls. + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("agent", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", + "gen_ai.usage.input_tokens", 15L, "gen_ai.usage.output_tokens", 7L))); + spans.add(span("chat1", "agent", "2026-01-01T00:00:02Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 10L, "gen_ai.usage.output_tokens", 4L))); + spans.add(span("chat2", "agent", "2026-01-01T00:00:03Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 5L, "gen_ai.usage.output_tokens", 3L))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + final Map root = byId(spans, "http").getAttributes(); + assertEquals(15L, ((Number) root.get("gen_ai.usage.input_tokens")).longValue()); + assertEquals(7L, ((Number) root.get("gen_ai.usage.output_tokens")).longValue()); + final Map agent = byId(spans, "agent").getAttributes(); + assertEquals(15L, ((Number) agent.get("gen_ai.usage.input_tokens")).longValue()); + } + + @Test + void countsANestedAgentsReportedTotalOnceOnTheMarkedSpan() { + // The outer agent reports no usage; a sub-agent reports the total of its own chat call. + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("agent", "http", "2026-01-01T00:00:01Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + spans.add(span("chat1", "agent", "2026-01-01T00:00:02Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 10L, "gen_ai.usage.output_tokens", 4L))); + spans.add(span("tool", "agent", "2026-01-01T00:00:03Z", Map.of("gen_ai.operation.name", "execute_tool"))); + spans.add(span("sub", "tool", "2026-01-01T00:00:04Z", Map.of( + "gen_ai.operation.name", "invoke_agent", + "gen_ai.usage.input_tokens", 5L, "gen_ai.usage.output_tokens", 3L))); + spans.add(span("subchat", "sub", "2026-01-01T00:00:05Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 5L, "gen_ai.usage.output_tokens", 3L))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agent"), markedSpanIds(spans)); + for (final String id : List.of("agent", "http")) { + final Map attrs = byId(spans, id).getAttributes(); + assertEquals(15L, ((Number) attrs.get("gen_ai.usage.input_tokens")).longValue(), id); + assertEquals(7L, ((Number) attrs.get("gen_ai.usage.output_tokens")).longValue(), id); + } + } + + @Test + void marksAGenAiRootItself() { + final List spans = new ArrayList<>(); + spans.add(span("agent", "", "2026-01-01T00:00:00Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + spans.add(span("chat1", "agent", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 10L, "gen_ai.usage.output_tokens", 4L))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agent"), markedSpanIds(spans)); + assertEquals(10L, ((Number) byId(spans, "agent").getAttributes().get("gen_ai.usage.input_tokens")).longValue()); + } + + @Test + void marksEachParallelTopLevelAgent() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("agentA", "http", "2026-01-01T00:00:01Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + spans.add(span("agentB", "http", "2026-01-01T00:00:02Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + spans.add(span("chatA", "agentA", "2026-01-01T00:00:03Z", Map.of("gen_ai.operation.name", "chat"))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agentA", "agentB"), markedSpanIds(spans)); + } + + @Test + void marksTheOutermostSpanInAnOrphanBatchWithoutTheRoot() { + // Children flushed after trace_flush_interval, before the root arrived. + final List spans = httpAgentTrace().stream() + .filter(s -> !"http".equals(s.getSpanId())) + .collect(Collectors.toCollection(ArrayList::new)); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agent"), markedSpanIds(spans)); + } + + @Test + void keepsAMarkerSetAtTheSource() { + final List spans = httpAgentTrace(); + // The SDK marked the HTTP handler span as the agent root. + spans.set(0, span("http", "", "2026-01-01T00:00:00Z", Map.of(DERIVED_GEN_AI_IS_ROOT_KEY, true))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("http"), markedSpanIds(spans)); + } + + @Test + void marksNothingWithoutGenAiSpans() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "GET"))); + spans.add(span("db", "http", "2026-01-01T00:00:01Z", Map.of("db.system", "postgresql"))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of(), markedSpanIds(spans)); + } + + @Test + void marksOnlyTheCallingAgentWhenAnAgentCallsAnotherAgentOverHttp() { + // POST /runs > invoke_agent Planner > HTTP client > HTTP server (other service) + // > invoke_agent Weather > MCP tool + final List spans = crossServiceTrace(); + Collections.shuffle(spans); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("planner"), markedSpanIds(spans)); + } + + @Test + void marksTheRemoteAgentWhenItsCallerIsMissingFromTheBatch() { + // The remote service's spans were flushed after trace_flush_interval without the caller's. + final List spans = crossServiceTrace().stream() + .filter(s -> List.of("server", "weather", "mcp").contains(s.getSpanId())) + .collect(Collectors.toCollection(ArrayList::new)); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("weather"), markedSpanIds(spans)); + } + + @Test + void doesNotMarkALateSpanWhoseAgentArrivedInAnEarlierBatch() { + final GenAiCoverage coverage = GenAiCoverage.create(1_000, java.time.Duration.ofMinutes(1)); + final List first = httpAgentTrace(); + final Span lateChat = first.remove(first.size() - 1); + + GenAiEnrichmentHelper.enrichBatch(first, coverage); + final List late = new ArrayList<>(List.of(lateChat)); + GenAiEnrichmentHelper.enrichBatch(late, coverage); + + assertEquals(List.of("agent"), markedSpanIds(first)); + assertEquals(List.of(), markedSpanIds(late)); + } + + @Test + void doesNotMarkALateRemoteAgentWhoseCallerArrivedEarlier() { + // The called service exports after the caller's trace has been processed. + final GenAiCoverage coverage = GenAiCoverage.create(1_000, java.time.Duration.ofMinutes(1)); + final List caller = crossServiceTrace().stream() + .filter(s -> List.of("http", "planner", "client").contains(s.getSpanId())) + .collect(Collectors.toCollection(ArrayList::new)); + final List remote = crossServiceTrace().stream() + .filter(s -> List.of("server", "weather", "mcp").contains(s.getSpanId())) + .collect(Collectors.toCollection(ArrayList::new)); + + GenAiEnrichmentHelper.enrichBatch(caller, coverage); + GenAiEnrichmentHelper.enrichBatch(remote, coverage); + + assertEquals(List.of("planner"), markedSpanIds(caller)); + assertEquals(List.of(), markedSpanIds(remote)); + } + + @Test + void doesNotMarkALateAgentBelowASpanMarkedAtTheSource() { + final GenAiCoverage coverage = GenAiCoverage.create(1_000, java.time.Duration.ofMinutes(1)); + final List first = new ArrayList<>(List.of( + span("http", "", "2026-01-01T00:00:00Z", Map.of(DERIVED_GEN_AI_IS_ROOT_KEY, true)))); + final List late = new ArrayList<>(List.of( + span("agent", "http", "2026-01-01T00:00:01Z", Map.of("gen_ai.operation.name", "invoke_agent")))); + + GenAiEnrichmentHelper.enrichBatch(first, coverage); + GenAiEnrichmentHelper.enrichBatch(late, coverage); + + assertEquals(List.of("http"), markedSpanIds(first)); + assertEquals(List.of(), markedSpanIds(late)); + } + + @Test + void doesNotMarkALateChildThroughTheProcessor() { + final OtelTraceRawProcessorConfig config = mock(OtelTraceRawProcessorConfig.class); + when(config.getTraceFlushIntervalSeconds()).thenReturn(180L); + when(config.getTraceGroupCacheMaxSize()).thenReturn(OtelTraceRawProcessorConfig.MAX_TRACE_ID_CACHE_SIZE); + when(config.getTraceGroupCacheTimeToLive()).thenReturn(OtelTraceRawProcessorConfig.DEFAULT_TRACE_ID_TTL); + final OTelTraceRawProcessor processor = + new OTelTraceRawProcessor(config, mock(PipelineDescription.class), mock(PluginMetrics.class)); + + final List spans = httpAgentTrace(); + final Span lateChat = spans.remove(spans.size() - 1); + final Span root = spans.remove(0); + processor.doExecute(records(spans)); + final List flushed = processor.doExecute(records(List.of(root))).stream() + .map(Record::getData).collect(Collectors.toList()); + // Exported after the root: emitted on its own, without its ancestors. + final List late = processor.doExecute(records(List.of(lateChat))).stream() + .map(Record::getData).collect(Collectors.toList()); + + assertEquals(List.of("agent"), markedSpanIds(flushed)); + assertEquals(1, late.size()); + assertEquals(List.of(), markedSpanIds(late)); + } + + @Test + void marksTheAgentSpanThroughTheProcessorWhenTheRootArrivesLast() { + final OtelTraceRawProcessorConfig config = mock(OtelTraceRawProcessorConfig.class); + when(config.getTraceFlushIntervalSeconds()).thenReturn(180L); + when(config.getTraceGroupCacheMaxSize()).thenReturn(OtelTraceRawProcessorConfig.MAX_TRACE_ID_CACHE_SIZE); + when(config.getTraceGroupCacheTimeToLive()).thenReturn(OtelTraceRawProcessorConfig.DEFAULT_TRACE_ID_TTL); + final OTelTraceRawProcessor processor = + new OTelTraceRawProcessor(config, mock(PipelineDescription.class), mock(PluginMetrics.class)); + + final List spans = httpAgentTrace(); + final Span root = spans.remove(0); + // Children are exported first and held until their root arrives. + assertEquals(0, processor.doExecute(records(spans)).size()); + final Collection> flushed = processor.doExecute(records(List.of(root))); + + final List out = flushed.stream().map(Record::getData).collect(Collectors.toList()); + assertEquals(List.of("agent"), markedSpanIds(out)); + assertFalse(byId(out, "http").getAttributes().containsKey(DERIVED_GEN_AI_IS_ROOT_KEY)); + } + + private static List crossServiceTrace() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("planner", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Planner"))); + spans.add(span("client", "planner", "2026-01-01T00:00:02Z", Map.of("http.method", "POST"))); + spans.add(span("server", "client", "2026-01-01T00:00:03Z", Map.of("http.route", "/weather"))); + spans.add(span("weather", "server", "2026-01-01T00:00:04Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Weather"))); + spans.add(span("mcp", "weather", "2026-01-01T00:00:05Z", Map.of( + "gen_ai.operation.name", "execute_tool"))); + return spans; + } + + /** POST /runs > invoke_agent > chat, execute_tool > chat. */ + @Test + void ordersAgentsByStartTimeWhateverTheFractionalPrecision() { + // As text, "...:01Z" sorts after "...:01.500Z"; as instants it is earlier. + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("late", "http", "2026-01-01T00:00:01.500Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Late Agent"))); + spans.add(span("early", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Early Agent"))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals("Early Agent", byId(spans, "http").getAttributes().get("gen_ai.agent.name")); + assertEquals(List.of("early", "late"), markedSpanIds(spans)); + } + + @Test + void keepsTheOutputTokensOfAnAgentThatReportsOnlyOutputTokens() { + final List spans = httpAgentTrace(); + spans.set(1, span("agent", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.usage.output_tokens", 700L))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + final Map agent = byId(spans, "agent").getAttributes(); + assertEquals(700L, ((Number) agent.get("gen_ai.usage.output_tokens")).longValue()); + assertNull(agent.get("gen_ai.usage.input_tokens")); + } + + @Test + void readsStringTokenCountsAndIgnoresNonNumericOnes() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("agent", "http", "2026-01-01T00:00:01Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + spans.add(span("chat1", "agent", "2026-01-01T00:00:02Z", Map.of( + "gen_ai.operation.name", "chat", "gen_ai.usage.input_tokens", "12", "gen_ai.usage.output_tokens", "n/a"))); + spans.add(span("chat2", "agent", "2026-01-01T00:00:03Z", Map.of( + "gen_ai.operation.name", "chat", "gen_ai.usage.input_tokens", 3L, "gen_ai.usage.output_tokens", 2L))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + final Map agent = byId(spans, "agent").getAttributes(); + assertEquals(15L, ((Number) agent.get("gen_ai.usage.input_tokens")).longValue()); + assertEquals(2L, ((Number) agent.get("gen_ai.usage.output_tokens")).longValue()); + } + + @Test + void neverMarksASpanSetToFalseAndStillMarksTheOthers() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST", DERIVED_GEN_AI_IS_ROOT_KEY, false))); + spans.add(span("agentA", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", DERIVED_GEN_AI_IS_ROOT_KEY, false))); + spans.add(span("agentB", "http", "2026-01-01T00:00:02Z", Map.of("gen_ai.operation.name", "invoke_agent"))); + + GenAiEnrichmentHelper.enrichBatch(spans); + + assertEquals(List.of("agentB"), markedSpanIds(spans)); + assertEquals(false, byId(spans, "agentA").getAttributes().get(DERIVED_GEN_AI_IS_ROOT_KEY)); + } + + @Test + void leavesTracesWithoutGenAiDataUntouched() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("db", "http", "2026-01-01T00:00:01Z", Map.of("db.system", "postgresql"))); + final GenAiCoverage coverage = GenAiCoverage.create(100, java.time.Duration.ofMinutes(1)); + + GenAiEnrichmentHelper.enrichBatch(spans, coverage); + + assertEquals(Map.of("http.method", "POST"), byId(spans, "http").getAttributes()); + assertEquals(0, coverage.estimatedSize()); + } + + private static List httpAgentTrace() { + final List spans = new ArrayList<>(); + spans.add(span("http", "", "2026-01-01T00:00:00Z", Map.of("http.method", "POST"))); + spans.add(span("agent", "http", "2026-01-01T00:00:01Z", Map.of( + "gen_ai.operation.name", "invoke_agent", "gen_ai.agent.name", "Travel Planner"))); + spans.add(span("chat1", "agent", "2026-01-01T00:00:02Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 10L, "gen_ai.usage.output_tokens", 4L))); + spans.add(span("tool", "agent", "2026-01-01T00:00:04Z", Map.of( + "gen_ai.operation.name", "execute_tool"))); + spans.add(span("chat2", "tool", "2026-01-01T00:00:05Z", Map.of( + "gen_ai.operation.name", "chat", + "gen_ai.usage.input_tokens", 5L, "gen_ai.usage.output_tokens", 3L))); + return spans; + } + + private static Span span(final String spanId, final String parentSpanId, final String startTime, + final Map attributes) { + return JacksonSpan.builder() + .withTraceId(TRACE_ID) + .withSpanId(spanId) + .withParentSpanId(parentSpanId) + .withTraceState("") + .withName(spanId) + .withKind("INTERNAL") + .withDurationInNanos(1L) + .withStartTime(startTime) + .withEndTime("2026-01-01T00:01:00Z") + .withTraceGroup(null) + .withTraceGroupFields(DefaultTraceGroupFields.builder().build()) + .withAttributes(new HashMap<>(attributes)) + .build(); + } + + private static List markedSpanIds(final Collection spans) { + return spans.stream() + .filter(s -> Boolean.TRUE.equals(s.getAttributes().get(DERIVED_GEN_AI_IS_ROOT_KEY))) + .map(Span::getSpanId) + .sorted() + .collect(Collectors.toList()); + } + + private static Span byId(final Collection spans, final String spanId) { + return spans.stream().filter(s -> spanId.equals(s.getSpanId())).findFirst().orElseThrow(); + } + + private static List> records(final List spans) { + return spans.stream().map(Record::new).collect(Collectors.toList()); + } +} diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverageTest.java b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverageTest.java new file mode 100644 index 0000000000..d7672d3f21 --- /dev/null +++ b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverageTest.java @@ -0,0 +1,92 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.oteltrace; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; + +class GenAiCoverageTest { + + @Test + void update_mergesTheReturnedSpanIdsWithTheKnownOnes() { + final GenAiCoverage coverage = GenAiCoverage.create(100, Duration.ofMinutes(1)); + + coverage.update("trace", known -> Set.of("a")); + coverage.update("trace", known -> { + assertThat(known, equalTo(Set.of("a"))); + return Set.of("b"); + }); + + assertThat(coverage.covered("trace"), equalTo(Set.of("a", "b"))); + assertThat(coverage.covered("other"), is(empty())); + } + + @Test + void update_onNone_runsWithNothingKnownAndKeepsNothing() { + GenAiCoverage.NONE.update("trace", known -> { + assertThat(known, is(empty())); + return Set.of("a"); + }); + + assertThat(GenAiCoverage.NONE.covered("trace"), is(empty())); + assertThat(GenAiCoverage.NONE.estimatedSize(), equalTo(0L)); + } + + @Test + void update_forOneTrace_runsOneEnrichmentAtATime() throws Exception { + final GenAiCoverage coverage = GenAiCoverage.create(100, Duration.ofMinutes(1)); + final CountDownLatch firstStarted = new CountDownLatch(1); + final CountDownLatch releaseFirst = new CountDownLatch(1); + final AtomicBoolean secondRanEarly = new AtomicBoolean(false); + final AtomicBoolean firstDone = new AtomicBoolean(false); + final ExecutorService executor = Executors.newFixedThreadPool(2); + try { + final Future first = executor.submit(() -> coverage.update("trace", known -> { + firstStarted.countDown(); + try { + releaseFirst.await(5, TimeUnit.SECONDS); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } + firstDone.set(true); + return Set.of("agent"); + })); + firstStarted.await(5, TimeUnit.SECONDS); + final Future second = executor.submit(() -> coverage.update("trace", known -> { + secondRanEarly.set(!firstDone.get()); + assertThat(known, equalTo(Set.of("agent"))); + return Set.of("child"); + })); + Thread.sleep(200); + releaseFirst.countDown(); + first.get(5, TimeUnit.SECONDS); + second.get(5, TimeUnit.SECONDS); + } finally { + executor.shutdownNow(); + } + + assertThat(secondRanEarly.get(), is(false)); + assertThat(coverage.covered("trace"), equalTo(Set.of("agent", "child"))); + } +} diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessorTest.java b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessorTest.java index 223e2d883b..707802d9e2 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessorTest.java +++ b/data-prepper-plugins/otel-trace-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessorTest.java @@ -391,6 +391,16 @@ void testMetricsOnTraceGroup() { assertThat(actualFunction.applyAsDouble(actualMeasuredObject), equalTo(2.0)); } + @Test + void testMetricsOnGenAiCoverage() { + final ArgumentCaptor gaugeObjectArgumentCaptor = ArgumentCaptor.forClass(Object.class); + final ArgumentCaptor gaugeFunctionArgumentCaptor = ArgumentCaptor.forClass(ToDoubleFunction.class); + verify(pluginMetrics).gauge(eq(OTelTraceRawProcessor.GENAI_COVERAGE_COUNT_METRIC_NAME), gaugeObjectArgumentCaptor.capture(), + gaugeFunctionArgumentCaptor.capture()); + + assertThat(gaugeFunctionArgumentCaptor.getValue().applyAsDouble(gaugeObjectArgumentCaptor.getValue()), equalTo(0.0)); + } + @Test void testMetricsOnSpanSet() { ArgumentCaptor gaugeObjectArgumentCaptor = ArgumentCaptor.forClass(Object.class);