llmgw/internal/router/trace.go

package router

import (
	"context"
	"net/http"
	"strings"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
)

// tracer is looked up per use rather than kept, so a provider installed after
// the router was built — as tests do — is the one spans go to.
func tracer() trace.Tracer { return otel.Tracer("gitea.mynodes.xyz/ilker/llmgw/internal/router") }

// startInference opens the span for one inference request, continuing the
// caller's trace when it sent a traceparent.
func startInference(req *http.Request, path, requestID string) (context.Context, trace.Span) {
	ctx := otel.GetTextMapPropagator().Extract(req.Context(), propagation.HeaderCarrier(req.Header))
	return tracer().Start(ctx, "inference "+path,
		trace.WithSpanKind(trace.SpanKindServer),
		trace.WithAttributes(
			attribute.String("llmgw.request_id", requestID),
			attribute.String("http.route", path),
		))
}

// endSpan closes a span, marking it failed when err is set.
func endSpan(span trace.Span, err error) {
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
	}
	span.End()
}

// withTraceparent is h with the current span's context in it, for the engine.
// A copy: the request's header is sent on every attempt, each its own span.
func withTraceparent(ctx context.Context, h http.Header) http.Header {
	if !trace.SpanContextFromContext(ctx).IsValid() {
		return h
	}
	out := h.Clone()
	if out == nil {
		out = http.Header{}
	}
	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(out))
	return out
}

// The OpenTelemetry GenAI semantic conventions: the names a trace backend
// reads an LLM call by, so llmgw's spans show in a customer's own tools —
// Grafana, Jaeger, Langfuse, anything that speaks OpenTelemetry — as model
// calls with their model and tokens, without a dashboard made for llmgw.
//
// What a request said and what the model answered are not among them, and
// are never added: llmgw keeps no prompt text anywhere, and a trace is the
// last place one should leak into.

// genAIOperation is the conventions' operation for an inference path, or ""
// for one they have no name for. The Anthropic Messages and Responses APIs
// are answered as chat completions, and so are chat.
func genAIOperation(path string) string {
	switch {
	case strings.HasSuffix(path, "/chat/completions"):
		return "chat"
	case strings.HasSuffix(path, "/completions"):
		return "text_completion"
	case strings.HasSuffix(path, "/embeddings"):
		return "embeddings"
	}
	return ""
}

// genAIRequest is what a request asked for, as the conventions name it.
func genAIRequest(path string, p *parsedRequest) []attribute.KeyValue {
	op := genAIOperation(path)
	if op == "" {
		return nil
	}
	kv := []attribute.KeyValue{
		attribute.String("gen_ai.operation.name", op),
		attribute.String("gen_ai.request.model", p.Model),
	}
	if p.MaxTokens > 0 {
		kv = append(kv, attribute.Int("gen_ai.request.max_tokens", p.MaxTokens))
	}
	if p.Temperature != nil {
		kv = append(kv, attribute.Float64("gen_ai.request.temperature", *p.Temperature))
	}
	if p.TopP != nil {
		kv = append(kv, attribute.Float64("gen_ai.request.top_p", *p.TopP))
	}
	if p.SessionID != "" {
		kv = append(kv, attribute.String("gen_ai.conversation.id", p.SessionID))
	}
	return kv
}

// genAIUsage is what an answer cost, in the conventions' names: the model
// that answered, and its tokens, cached ones apart.
func genAIUsage(path, served string, u *proxy.Usage) []attribute.KeyValue {
	if genAIOperation(path) == "" {
		return nil
	}
	var kv []attribute.KeyValue
	if served != "" {
		kv = append(kv, attribute.String("gen_ai.response.model", served))
	}
	if u != nil {
		kv = append(kv,
			attribute.Int("gen_ai.usage.input_tokens", u.PromptTokens),
			attribute.Int("gen_ai.usage.output_tokens", u.CompletionTokens))
		if u.CachedTokens > 0 {
			kv = append(kv, attribute.Int("gen_ai.usage.cache_read.input_tokens", u.CachedTokens))
		}
	}
	return kv
}

// genAIProvider names who answered an attempt: the engine serving a replica
// (vllm, llamacpp, …), or the outside provider's, for an external model.
func genAIProvider(engine string) string {
	if engine == "" {
		return "llmgw"
	}
	return engine
}