chore(llmobs): add telemetry to track llmobs usage (#4088)

Co-authored-by: darccio <[email protected]>

Author: rarguelloF

instrumentation/testutils/testutils.go

@@ -16,6 +16,7 @@ import (
 	"github.com/DataDog/dd-trace-go/v2/internal/globalconfig"
 	"github.com/DataDog/dd-trace-go/v2/internal/normalizer"
 	"github.com/DataDog/dd-trace-go/v2/internal/statsdtest"
+	"github.com/DataDog/dd-trace-go/v2/internal/telemetry"
 )
 
 func SetGlobalServiceName(t *testing.T, val string) {
@@ -116,3 +117,10 @@ func SetPropagatingTag(t testing.TB, ctx *tracer.SpanContext, k, v string) {
 	cc := (*cookieCutter)(*(*unsafe.Pointer)(unsafe.Pointer(&ptr)))
 	cc.trace.propagatingTags[k] = v
 }
+
+// FlushTelemetry flushes any pending telemetry data.
+func FlushTelemetry() {
+	if client := telemetry.GlobalClient(); client != nil {
+		client.Flush()
+	}
+}

internal/llmobs/llmobs.go

@@ -37,6 +37,12 @@ var (
 	errAgentlessRequiresAPIKey = errors.New("LLMOBs agentless mode requires a valid API key - set the DD_API_KEY env variable to configure one")
 	errMLAppRequired           = errors.New("ML App is required for sending LLM Observability data")
 	errAgentModeNotSupported   = errors.New("DD_LLMOBS_AGENTLESS_ENABLED has been configured to false but the agent is not available or does not support LLMObs")
+	errInvalidMetricLabel      = errors.New("label is required for evaluation metrics")
+	errFinishedSpan            = errors.New("span is already finished")
+	errEvalJoinBothPresent     = errors.New("provide either span/trace IDs or tag key/value, not both")
+	errEvalJoinNonePresent     = errors.New("must provide either span/trace IDs or tag key/value for joining")
+	errInvalidSpanJoin         = errors.New("both span and trace IDs are required for span-based joining")
+	errInvalidTagJoin          = errors.New("both tag key and value are required for tag-based joining")
 )
 
 const (
@@ -182,7 +188,11 @@ func newLLMObs(cfg *config.Config, tracer Tracer) (*LLMObs, error) {
 
 // Start starts the global LLMObs instance with the given configuration and tracer.
 // Returns an error if LLMObs is already running or if configuration is invalid.
-func Start(cfg config.Config, tracer Tracer) error {
+func Start(cfg config.Config, tracer Tracer) (err error) {
+	startTime := time.Now()
+	defer func() {
+		trackLLMObsStart(startTime, err, cfg)
+	}()
 	mu.Lock()
 	defer mu.Unlock()
 
@@ -224,6 +234,7 @@ func ActiveLLMObs() (*LLMObs, error) {
 func Flush() {
 	if activeLLMObs != nil {
 		activeLLMObs.Flush()
+		trackUserFlush()
 	}
 }
 
@@ -383,6 +394,7 @@ func (l *LLMObs) batchSend(params batchSendParams) {
 			}
 			if err := l.Transport.PushSpanEvents(ctx, events); err != nil {
 				log.Error("llmobs: failed to push span events: %v", err.Error())
+				trackDroppedPayload(len(events), telemetryMetricDroppedSpanEvents, "transport_error")
 			} else {
 				log.Debug("llmobs: push span events success")
 			}
@@ -403,6 +415,7 @@ func (l *LLMObs) batchSend(params batchSendParams) {
 			}
 			if err := l.Transport.PushEvalMetrics(ctx, metrics); err != nil {
 				log.Error("llmobs: failed to push eval metrics: %v", err.Error())
+				trackDroppedPayload(len(metrics), telemetryMetricDroppedEvalEvents, "transport_error")
 			} else {
 				log.Debug("llmobs: push eval metrics success")
 			}
@@ -583,20 +596,34 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 		Scope:            span.scope,
 	}
 	if b, err := json.Marshal(ev); err == nil {
-		if len(b) > sizeLimitEVPEvent {
+		rawSize := len(b)
+		trackSpanEventRawSize(ev, rawSize)
+
+		truncated := false
+		if rawSize > sizeLimitEVPEvent {
 			log.Warn(
 				"llmobs: dropping llmobs span event input/output because its size (%s) exceeds the event size limit (5MB)",
-				readableBytes(len(b)),
+				readableBytes(rawSize),
 			)
-			dropSpanEventIO(ev)
+			truncated = dropSpanEventIO(ev)
+			if !truncated {
+				log.Debug("llmobs: attempted to drop span event IO but it was not present")
+			}
+		}
+		actualSize := rawSize
+		if truncated {
+			if b, err := json.Marshal(ev); err == nil {
+				actualSize = len(b)
+			}
 		}
+		trackSpanEventSize(ev, actualSize, truncated)
 	}
 	return ev
 }
 
-func dropSpanEventIO(ev *transport.LLMObsSpanEvent) {
+func dropSpanEventIO(ev *transport.LLMObsSpanEvent) bool {
 	if ev == nil {
-		return
+		return false
 	}
 	droppedIO := false
 	if _, ok := ev.Meta["input"]; ok {
@@ -612,11 +639,14 @@ func dropSpanEventIO(ev *transport.LLMObsSpanEvent) {
 	} else {
 		log.Debug("llmobs: attempted to drop span event IO but it was not present")
 	}
+	return droppedIO
 }
 
 // StartSpan starts a new LLMObs span with the given kind, name, and configuration.
 // Returns the created span and a context containing the span.
 func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg StartSpanConfig) (*Span, context.Context) {
+	defer trackSpanStarted()
+
 	spanName := name
 	if spanName == "" {
 		spanName = string(kind)
@@ -692,20 +722,38 @@ func (l *LLMObs) StartExperimentSpan(ctx context.Context, name string, experimen
 
 // SubmitEvaluation submits an evaluation metric for a span.
 // The span can be identified either by span/trace IDs or by tag key-value pairs.
-func (l *LLMObs) SubmitEvaluation(cfg EvaluationConfig) error {
-	// Validate exactly one join method is provided
-	hasSpanJoin := cfg.SpanID != "" && cfg.TraceID != ""
-	hasTagJoin := cfg.TagKey != "" && cfg.TagValue != ""
+func (l *LLMObs) SubmitEvaluation(cfg EvaluationConfig) (err error) {
+	var metric *transport.LLMObsMetric
+	defer func() {
+		trackSubmitEvaluationMetric(metric, err)
+	}()
 
+	if cfg.Label == "" {
+		return errInvalidMetricLabel
+	}
+	var (
+		hasTagJoin  bool
+		hasSpanJoin bool
+	)
+	if cfg.SpanID != "" || cfg.TraceID != "" {
+		if !(cfg.SpanID != "" && cfg.TraceID != "") {
+			return errInvalidSpanJoin
+		}
+		hasSpanJoin = true
+	}
+	if cfg.TagKey != "" || cfg.TagValue != "" {
+		if !(cfg.TagKey != "" && cfg.TagValue != "") {
+			return errInvalidTagJoin
+		}
+		hasTagJoin = true
+	}
 	if hasSpanJoin && hasTagJoin {
-		return errors.New("provide either span/trace IDs or tag key/value, not both")
+		return errEvalJoinBothPresent
 	}
 	if !hasSpanJoin && !hasTagJoin {
-		return errors.New("must provide either span/trace IDs or tag key/value for joining")
-	}
-	if cfg.Label == "" {
-		return errors.New("label is required for evaluation metrics")
+		return errEvalJoinNonePresent
 	}
+
 	numValues := 0
 	if cfg.CategoricalValue != nil {
 		numValues++
@@ -751,7 +799,7 @@ func (l *LLMObs) SubmitEvaluation(cfg EvaluationConfig) error {
 	}
 	tags = append(tags, fmt.Sprintf("ddtrace.version:%s", version.Tag))
 
-	metric := &transport.LLMObsMetric{
+	metric = &transport.LLMObsMetric{
 		JoinOn:      joinOn,
 		Label:       cfg.Label,
 		MLApp:       mlApp,

internal/llmobs/llmobs_test.go

@@ -1704,7 +1704,7 @@ func TestLLMObsLifecycle(t *testing.T) {
 	})
 }
 
-func BenchmarkStartSpan(b *testing.B) {
+func BenchmarkLLMObsStartSpan(b *testing.B) {
 	run := func(b *testing.B, ll *llmobs.LLMObs, tt *testtracer.TestTracer, done chan struct{}) {
 		b.Log("starting benchmark")
 

internal/llmobs/span.go

@@ -283,6 +283,9 @@ func (s *Span) Finish(cfg FinishSpanConfig) {
 		log.Debug("llmobs: attempted to finish an already finished span")
 		return
 	}
+	defer func() {
+		trackSpanFinished(s)
+	}()
 
 	if cfg.FinishTime.IsZero() {
 		cfg.FinishTime = time.Now()
@@ -303,17 +306,23 @@ func (s *Span) Finish(cfg FinishSpanConfig) {
 	}
 	l.submitLLMObsSpan(s)
 	s.finished = true
-
-	//TODO: telemetry.record_span_created(span)
 }
 
 // Annotate adds annotations to the span using the provided SpanAnnotations.
 func (s *Span) Annotate(a SpanAnnotations) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	var err error
+	defer func() {
+		if err != nil {
+			log.Warn("llmobs: failed to annotate span: %v", err.Error())
+		}
+		trackSpanAnnotations(s, err)
+	}()
+
 	if s.finished {
-		log.Warn("llmobs: cannot annotate a finished span")
+		err = errFinishedSpan
 		return
 	}