72 lines
1.9 KiB
Go
72 lines
1.9 KiB
Go
package litellm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
|
|
"go.opentelemetry.io/otel/sdk/resource"
|
|
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
|
semconv "go.opentelemetry.io/otel/semconv/v1.36.0"
|
|
"google.golang.org/adk/telemetry"
|
|
)
|
|
|
|
// SetupTelemetry wires ADK OTLP tracing from environment variables.
|
|
//
|
|
// Reads:
|
|
// - OTLP_ENDPOINT full URL base, e.g. http://jaeger.d-ma.be:4318 (skip if empty)
|
|
// - ADK_SERVICE_NAME service name in Jaeger (default: "agent")
|
|
// - ADK_SERVICE_VERSION semver label (default: "0.1.0")
|
|
//
|
|
// Returns a shutdown func to call on exit with a short-timeout context.
|
|
// No-op (nil error, noop shutdown) when OTLP_ENDPOINT is unset.
|
|
func SetupTelemetry(ctx context.Context) (shutdown func(context.Context) error, err error) {
|
|
endpoint := os.Getenv("OTLP_ENDPOINT")
|
|
if endpoint == "" {
|
|
return func(context.Context) error { return nil }, nil
|
|
}
|
|
|
|
svcName := os.Getenv("ADK_SERVICE_NAME")
|
|
if svcName == "" {
|
|
svcName = "agent"
|
|
}
|
|
svcVersion := os.Getenv("ADK_SERVICE_VERSION")
|
|
if svcVersion == "" {
|
|
svcVersion = "0.1.0"
|
|
}
|
|
|
|
exporter, err := otlptracehttp.New(ctx,
|
|
otlptracehttp.WithEndpointURL(endpoint+"/v1/traces"),
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("otlp exporter: %w", err)
|
|
}
|
|
|
|
res, err := resource.New(ctx,
|
|
resource.WithAttributes(
|
|
semconv.ServiceName(svcName),
|
|
semconv.ServiceVersion(svcVersion),
|
|
),
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("resource: %w", err)
|
|
}
|
|
|
|
providers, err := telemetry.New(ctx,
|
|
telemetry.WithSpanProcessors(sdktrace.NewBatchSpanProcessor(exporter)),
|
|
telemetry.WithResource(res),
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("telemetry.New: %w", err)
|
|
}
|
|
providers.SetGlobalOtelProviders()
|
|
|
|
return func(ctx context.Context) error {
|
|
shutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
|
defer cancel()
|
|
return providers.Shutdown(shutCtx)
|
|
}, nil
|
|
}
|