llmgw/internal/router/trace_test.go

package router

import (
	"strings"
	"testing"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.opentelemetry.io/otel/trace/noop"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)

func TestARequestIsTracedIntoTheEngine(t *testing.T) {
	rec := tracetest.NewSpanRecorder()
	prevTP, prevProp := otel.GetTracerProvider(), otel.GetTextMapPropagator()
	otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)))
	otel.SetTextMapPropagator(propagation.TraceContext{})
	t.Cleanup(func() { otel.SetTracerProvider(prevTP); otel.SetTextMapPropagator(prevProp) })

	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	// The caller's own trace, continued rather than replaced.
	const caller = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
	res := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}},
		map[string]string{"traceparent": caller, "X-Request-Id": "req-1"})
	if res.Code != 200 {
		t.Fatalf("status %d: %s", res.Code, res.Body.String())
	}

	spans := map[string]sdktrace.ReadOnlySpan{}
	for _, s := range rec.Ended() {
		spans[s.Name()] = s
	}
	root, ok := spans["inference /v1/chat/completions"]
	if !ok {
		t.Fatalf("no inference span among %v", spans)
	}
	if got := root.SpanContext().TraceID().String(); got != "4bf92f3577b34da6a3ce929d0e0e4736" {
		t.Errorf("the caller's trace was not continued: %s", got)
	}
	attrs := map[string]string{}
	for _, a := range root.Attributes() {
		attrs[string(a.Key)] = a.Value.Emit()
	}
	if attrs["llmgw.request_id"] != "req-1" || attrs["llmgw.model"] != "m" ||
		attrs["http.response.status_code"] != "200" || attrs["llmgw.endpoint"] != "a" {
		t.Errorf("inference span attributes %v", attrs)
	}
	for _, name := range []string{"admission", "schedule", "upstream"} {
		s, ok := spans[name]
		if !ok {
			t.Errorf("no %s span", name)
			continue
		}
		if s.Parent().SpanID() != root.SpanContext().SpanID() {
			t.Errorf("%s is not under the request's span", name)
		}
	}
	up := spans["upstream"]
	tp, _ := e.traceparent.Load().(string)
	if want := "00-" + up.SpanContext().TraceID().String() + "-" + up.SpanContext().SpanID().String() + "-01"; tp != want {
		t.Errorf("the engine got traceparent %q, expected %q", tp, want)
	}
}

func TestWithoutTracingTheEngineGetsNoTraceparent(t *testing.T) {
	// What an unconfigured process has. Set explicitly, since the default
	// global delegates for good to the first real provider another test set.
	prevTP, prevProp := otel.GetTracerProvider(), otel.GetTextMapPropagator()
	otel.SetTracerProvider(noop.NewTracerProvider())
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator())
	t.Cleanup(func() { otel.SetTracerProvider(prevTP); otel.SetTextMapPropagator(prevProp) })
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	if res := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil); res.Code != 200 {
		t.Fatalf("status %d", res.Code)
	}
	if tp, _ := e.traceparent.Load().(string); tp != "" {
		t.Errorf("traceparent %q sent with tracing off", tp)
	}
}

func TestTracesSpeakTheGenAIConventionsAndHoldNoText(t *testing.T) {
	rec := tracetest.NewSpanRecorder()
	prevTP := otel.GetTracerProvider()
	otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)))
	t.Cleanup(func() { otel.SetTracerProvider(prevTP) })

	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	res := post(t, rt.Handler(), map[string]any{"model": "m", "max_tokens": 64, "temperature": 0.2, "top_p": 0.9,
		"messages": []any{map[string]any{"role": "user", "content": "the secret word is aubergine"}}}, nil)
	if res.Code != 200 {
		t.Fatalf("status %d: %s", res.Code, res.Body.String())
	}
	got := map[string]map[string]string{}
	for _, s := range rec.Ended() {
		attrs := map[string]string{}
		for _, a := range s.Attributes() {
			attrs[string(a.Key)] = a.Value.Emit()
			// Neither the prompt nor the reply ("hi"), in any attribute.
			if v := a.Value.Emit(); strings.Contains(v, "aubergine") || v == "hi" {
				t.Errorf("span %s carries text in %s", s.Name(), a.Key)
			}
		}
		got[s.Name()] = attrs
	}
	root := got["inference /v1/chat/completions"]
	for k, want := range map[string]string{
		"gen_ai.operation.name":      "chat",
		"gen_ai.request.model":       "m",
		"gen_ai.request.max_tokens":  "64",
		"gen_ai.request.temperature": "0.2",
		"gen_ai.request.top_p":       "0.9",
		"gen_ai.response.model":      "m",
		"gen_ai.usage.input_tokens":  "4",
		"gen_ai.usage.output_tokens": "2",
	} {
		if root[k] != want {
			t.Errorf("request span %s = %q, want %q", k, root[k], want)
		}
	}
	up := got["upstream"]
	if up["gen_ai.operation.name"] != "chat" || up["gen_ai.provider.name"] == "" || up["gen_ai.usage.output_tokens"] != "2" {
		t.Errorf("upstream span %v", up)
	}
}

func TestEveryInferencePathHasItsGenAIOperation(t *testing.T) {
	for path, want := range map[string]string{
		"/v1/chat/completions": "chat",
		"/v1/completions":      "text_completion",
		"/v1/embeddings":       "embeddings",
		"/v1/rerank":           "",
	} {
		if got := genAIOperation(path); got != want {
			t.Errorf("%s: %q, want %q", path, got, want)
		}
	}
}

func TestATemperatureWrittenAsAStringIsTheEnginesToJudge(t *testing.T) {
	// Read for the trace only: before it was, "0.7" reached the engine.
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	rec := post(t, rt.Handler(), map[string]any{"model": "m", "temperature": "0.7",
		"messages": []any{map[string]any{"role": "user", "content": "hi"}}}, nil)
	if rec.Code != 200 {
		t.Errorf("status %d: %s", rec.Code, rec.Body.String())
	}
}