wip

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

cmd/dynamic_module/main.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/envoyproxy/ai-gateway/internal/tracing"
 	"github.com/prometheus/client_golang/prometheus"
 	otelprom "go.opentelemetry.io/otel/exporters/prometheus"
 
@@ -99,6 +100,17 @@ func (g *globalState) initializeEnv() error {
 	if err != nil {
 		return fmt.Errorf("failed to parse metrics header mapping: %w", err)
 	}
+	spanRequestHeaderAttributes, err := internalapi.ParseRequestHeaderAttributeMapping(os.Getenv(
+		"AI_GATEWAY_DYNAMIC_MODULE_FILTER_TRACING_REQUEST_HEADER_ATTRIBUTES",
+	))
+	if err != nil {
+		return fmt.Errorf("failed to parse tracing header mapping: %w", err)
+	}
+
+	tracing, err := tracing.NewTracingFromEnv(ctx, os.Stdout, spanRequestHeaderAttributes)
+	if err != nil {
+		return err
+	}
 
 	g.env = &dynamicmodule.Env{
 		RootPrefix:                    os.Getenv("AI_GATEWAY_DYNAMIC_MODULE_ROOT_PREFIX"),
@@ -109,6 +121,7 @@ func (g *globalState) initializeEnv() error {
 		EmbeddingsMetricsFactory:      metrics.NewMetricsFactory(meter, metricsRequestHeaderAttributes, metrics.GenAIOperationEmbedding),
 		ImageGenerationMetricsFactory: metrics.NewMetricsFactory(meter, metricsRequestHeaderAttributes, metrics.GenAIOperationImageGeneration),
 		RerankMetricsFactory:          metrics.NewMetricsFactory(meter, metricsRequestHeaderAttributes, metrics.GenAIOperationRerank),
+		Tracing:                       tracing,
 	}
 	return nil
 }

internal/dynamicmodule/dynamic_module.go

@@ -8,6 +8,7 @@ package dynamicmodule
 import (
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 	"github.com/envoyproxy/ai-gateway/internal/metrics"
+	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
 )
 
 // endpoint represents the type of the endpoint that the request is targeting.
@@ -40,4 +41,5 @@ type Env struct {
 	EmbeddingsMetricsFactory,
 	ImageGenerationMetricsFactory,
 	RerankMetricsFactory metrics.Factory
+	Tracing tracing.Tracing
 }

internal/dynamicmodule/router_filter.go

@@ -6,21 +6,22 @@
 package dynamicmodule
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"path"
 	"strings"
 	"unsafe"
 
-	openaisdk "github.com/openai/openai-go/v2"
-
 	"github.com/envoyproxy/ai-gateway/internal/apischema/anthropic"
 	cohereschema "github.com/envoyproxy/ai-gateway/internal/apischema/cohere"
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	"github.com/envoyproxy/ai-gateway/internal/dynamicmodule/sdk"
+	"github.com/envoyproxy/ai-gateway/internal/endpointspec"
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
+	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
 )
 
 const routerFilterPointerDynamicMetadataKey = "router_filter_pointer"
@@ -33,17 +34,41 @@ type (
 	routerFilterConfig struct {
 		fcr              **filterapi.RuntimeConfig
 		prefixToEndpoint map[string]endpoint
+		tracing          tracing.Tracing
 	}
 	// routerFilter implements [sdk.HTTPFilter].
 	routerFilter struct {
-		routerFilterConfig     *routerFilterConfig
+		// prefixToEndpoint maps request path prefixes to endpoints. Shallow copy of
+		// the one in routerFilterConfig at the time of filter creation.
+		prefixToEndpoint map[string]endpoint
+		// runtimeFilterConfig is the snapshot of the runtime filter configuration at the time of filter creation.
+		runtimeFilterConfig *filterapi.RuntimeConfig
+		// tracing is the tracing implementation inherited from the environment.
+		tracing tracing.Tracing
+
+		// endpoint is the endpoint that the current request is targeting.
+		endpoint endpoint
+		// typedFilter is the typed router filter for the current request.
+		typedFilter routerFilterTypedIface
+	}
+
+	// routerFilterTypedIface is the interface for the typed router filter.
+	routerFilterTypedIface interface {
+		RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.RequestBodyStatus
+	}
+
+	// routerFilter typed is the typed implementation of the router filter for a specific endpoint.
+	routerFilterTyped[ReqT, RespT, RespChunkT any, EndpointSpec endpointspec.Spec[ReqT, RespT, RespChunkT]] struct {
 		runtimeFilterConfig    *filterapi.RuntimeConfig
-		endpoint               endpoint
-		originalHeaders        map[string]string
-		originalRequestBody    any
-		originalRequestBodyRaw []byte
-		span                   any
 		attemptCount           int
+		ep                     EndpointSpec
+		originalRequestHeaders map[string]string
+		originalRequestBody    *ReqT
+		originalRequestBodyRaw []byte
+		originalModel          internalapi.OriginalModel
+		stream                 bool
+		tracer                 tracing.RequestTracer[ReqT, RespT, RespChunkT]
+		span                   tracing.Span[RespT, RespChunkT]
 	}
 )
 
@@ -61,12 +86,17 @@ func NewRouterFilterConfig(env *Env, fcr **filterapi.RuntimeConfig) sdk.HTTPFilt
 	return &routerFilterConfig{
 		fcr:              fcr,
 		prefixToEndpoint: prefixToEndpoint,
+		tracing:          env.Tracing,
 	}
 }
 
 // NewFilter implements [sdk.HTTPFilterConfig].
 func (f *routerFilterConfig) NewFilter() sdk.HTTPFilter {
-	return &routerFilter{routerFilterConfig: f, runtimeFilterConfig: *f.fcr}
+	return &routerFilter{
+		prefixToEndpoint:    f.prefixToEndpoint,
+		runtimeFilterConfig: *f.fcr,
+		tracing:             f.tracing,
+	}
 }
 
 // RequestHeaders implements [sdk.HTTPFilter].
@@ -76,7 +106,7 @@ func (f *routerFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Request
 	if queryIndex := strings.Index(p, "?"); queryIndex != -1 {
 		p = p[:queryIndex]
 	}
-	ep, ok := f.routerFilterConfig.prefixToEndpoint[p]
+	ep, ok := f.prefixToEndpoint[p]
 	if !ok {
 		e.SendLocalReply(404, nil, []byte(fmt.Sprintf("unsupported path: %s", p)))
 		return sdk.RequestHeadersStatusStopIteration
@@ -90,6 +120,45 @@ func (f *routerFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Request
 
 // RequestBody implements [sdk.HTTPFilter].
 func (f *routerFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.RequestBodyStatus {
+	switch f.endpoint {
+	case chatCompletionsEndpoint:
+		f.typedFilter = &routerFilterTyped[openai.ChatCompletionRequest, openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk, endpointspec.ChatCompletionsEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.ChatCompletionTracer(),
+		}
+	case completionsEndpoint:
+		f.typedFilter = &routerFilterTyped[openai.CompletionRequest, openai.CompletionResponse, openai.CompletionResponse, endpointspec.CompletionsEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.CompletionTracer(),
+		}
+	case embeddingsEndpoint:
+		f.typedFilter = &routerFilterTyped[openai.EmbeddingRequest, openai.EmbeddingResponse, struct{}, endpointspec.EmbeddingsEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.EmbeddingsTracer(),
+		}
+	case imagesGenerationsEndpoint:
+		f.typedFilter = &routerFilterTyped[openai.ImageGenerationRequest, openai.ImageGenerationResponse, struct{}, endpointspec.ImageGenerationEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.ImageGenerationTracer(),
+		}
+	case rerankEndpoint:
+		f.typedFilter = &routerFilterTyped[cohereschema.RerankV2Request, cohereschema.RerankV2Response, struct{}, endpointspec.RerankEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.RerankTracer(),
+		}
+	case messagesEndpoint:
+		f.typedFilter = &routerFilterTyped[anthropic.MessagesRequest, anthropic.MessagesResponse, anthropic.MessagesStreamChunk, endpointspec.MessagesEndpointSpec]{
+			runtimeFilterConfig: f.runtimeFilterConfig,
+			tracer:              f.tracing.MessageTracer(),
+		}
+	default:
+		e.SendLocalReply(500, nil, []byte("BUG: unsupported endpoint at body parsing: "+fmt.Sprintf("%d", f.endpoint)))
+		return sdk.RequestBodyStatusStopIterationAndBuffer
+	}
+	return f.typedFilter.RequestBody(e, endOfStream)
+}
+
+func (f *routerFilterTyped[ReqT, RespT, RespChunkT, EndpointSpecT]) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.RequestBodyStatus {
 	if !endOfStream {
 		return sdk.RequestBodyStatusStopIterationAndBuffer
 	}
@@ -103,39 +172,31 @@ func (f *routerFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.
 		e.SendLocalReply(400, nil, []byte("failed to read request body: "+err.Error()))
 		return sdk.RequestBodyStatusStopIterationAndBuffer
 	}
+
 	f.originalRequestBodyRaw = raw
-	var parsed any
-	var modelName string
-	switch f.endpoint {
-	case chatCompletionsEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.ChatCompletionRequest) string { return req.Model })
-	case completionsEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.CompletionRequest) string { return req.Model })
-	case embeddingsEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.EmbeddingRequest) string { return req.Model })
-	case imagesGenerationsEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *openaisdk.ImageGenerateParams) string { return req.Model })
-	case rerankEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *cohereschema.RerankV2Request) string { return req.Model })
-	case messagesEndpoint:
-		parsed, modelName, err = parseBodyWithModel(raw, func(req *anthropic.MessagesRequest) string { return req.GetModel() })
-	default:
-		e.SendLocalReply(500, nil, []byte("BUG: unsupported endpoint at body parsing: "+fmt.Sprintf("%d", f.endpoint)))
-	}
-	if err != nil {
-		e.SendLocalReply(400, nil, []byte("failed to parse request body: "+err.Error()))
-		return sdk.RequestBodyStatusStopIterationAndBuffer
+	var maybeMutatedOriginalBodyRaw []byte
+	f.originalModel, f.originalRequestBody, f.stream, maybeMutatedOriginalBodyRaw, err =
+		f.ep.ParseBody(raw, len(f.runtimeFilterConfig.RequestCosts) > 0)
+	if len(maybeMutatedOriginalBodyRaw) > 0 {
+		f.originalRequestBodyRaw = maybeMutatedOriginalBodyRaw
 	}
-	f.originalRequestBody = parsed
-	if !e.SetRequestHeader(internalapi.ModelNameHeaderKeyDefault, []byte(modelName)) {
+	if !e.SetRequestHeader(internalapi.ModelNameHeaderKeyDefault, []byte(f.originalModel)) {
 		e.SendLocalReply(500, nil, []byte("failed to set model name header"))
 		return sdk.RequestBodyStatusStopIterationAndBuffer
 	}
 	// Store the pointer to the filter in dynamic metadata for later retrieval in the upstream filter.
 	e.SetDynamicMetadataString(internalapi.AIGatewayFilterMetadataNamespace, routerFilterPointerDynamicMetadataKey,
 		fmt.Sprintf("%d", uintptr(unsafe.Pointer(f))))
 
-	f.originalHeaders = multiValueHeadersToSingleValue(e.GetRequestHeaders())
+	f.originalRequestHeaders = multiValueHeadersToSingleValue(e.GetRequestHeaders())
+
+	f.span = f.tracer.StartSpanAndInjectHeaders(
+		context.Background(),
+		f.originalRequestHeaders,
+		&headerMutationCarrier{e: e},
+		f.originalRequestBody,
+		f.originalRequestBodyRaw,
+	)
 	return sdk.RequestBodyStatusContinue
 }
 
@@ -174,14 +235,6 @@ func (f *routerFilter) handleModelsEndpoint(e sdk.EnvoyHTTPFilter) sdk.RequestHe
 	return sdk.RequestHeadersStatusStopIteration
 }
 
-func parseBodyWithModel[T any](body []byte, modelExtractFn func(req *T) string) (interface{}, string, error) {
-	var req T
-	if err := json.Unmarshal(body, &req); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return req, modelExtractFn(&req), nil
-}
-
 // multiValueHeadersToSingleValue converts a map of headers with multiple values to a map of headers with single values by taking the first value for each header.
 //
 // TODO: this is purely for feature parity with the old filter where we ignore the case of multiple header values.
@@ -192,3 +245,23 @@ func multiValueHeadersToSingleValue(headers map[string][]string) map[string]stri
 	}
 	return singleValueHeaders
 }
+
+// headerMutationCarrier implements [propagation.TextMapCarrier].
+type headerMutationCarrier struct {
+	e sdk.EnvoyHTTPFilter
+}
+
+// Get implements the same method as defined on propagation.TextMapCarrier.
+func (c *headerMutationCarrier) Get(string) string {
+	panic("unexpected as this carrier is write-only for injection")
+}
+
+// Set adds a key-value pair to the HeaderMutation.
+func (c *headerMutationCarrier) Set(key, value string) {
+	_ = c.e.SetResponseHeader(key, []byte(value))
+}
+
+// Keys implements the same method as defined on propagation.TextMapCarrier.
+func (c *headerMutationCarrier) Keys() []string {
+	panic("unexpected as this carrier is write-only for injection")
+}

internal/dynamicmodule/upstream_filter.go

@@ -194,71 +194,3 @@ func (f *upstreamFilter) ResponseHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Resp
 func (f *upstreamFilter) ResponseBody(sdk.EnvoyHTTPFilter, bool) sdk.ResponseBodyStatus {
 	return sdk.ResponseBodyStatusContinue
 }
-
-func (f *upstreamFilter) initializeTranslatorMetrics(b *filterapi.Backend) error {
-	out := b.Schema
-	modelNameOverride := b.ModelNameOverride
-	switch f.rf.endpoint {
-	case chatCompletionsEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaOpenAI:
-			f.translator = translator.NewChatCompletionOpenAIToOpenAITranslator(out.Version, modelNameOverride)
-		case filterapi.APISchemaAWSBedrock:
-			f.translator = translator.NewChatCompletionOpenAIToAWSBedrockTranslator(modelNameOverride)
-		case filterapi.APISchemaAzureOpenAI:
-			f.translator = translator.NewChatCompletionOpenAIToAzureOpenAITranslator(out.Version, modelNameOverride)
-		case filterapi.APISchemaGCPVertexAI:
-			f.translator = translator.NewChatCompletionOpenAIToGCPVertexAITranslator(modelNameOverride)
-		case filterapi.APISchemaGCPAnthropic:
-			f.translator = translator.NewChatCompletionOpenAIToGCPAnthropicTranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.ChatCompletionMetricsFactory.NewMetrics()
-	case completionsEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaOpenAI:
-			f.translator = translator.NewChatCompletionOpenAIToOpenAITranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.CompletionMetricsFactory.NewMetrics()
-	case embeddingsEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaOpenAI:
-			f.translator = translator.NewEmbeddingOpenAIToOpenAITranslator(out.Version, modelNameOverride)
-		case filterapi.APISchemaAzureOpenAI:
-			f.translator = translator.NewEmbeddingOpenAIToAzureOpenAITranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.CompletionMetricsFactory.NewMetrics()
-	case imagesGenerationsEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaOpenAI:
-			f.translator = translator.NewImageGenerationOpenAIToOpenAITranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.CompletionMetricsFactory.NewMetrics()
-	case rerankEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaCohere:
-			f.translator = translator.NewRerankCohereToCohereTranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.RerankMetricsFactory.NewMetrics()
-	case messagesEndpoint:
-		switch out.Name {
-		case filterapi.APISchemaAnthropic:
-			f.translator = translator.NewAnthropicToAnthropicTranslator(out.Version, modelNameOverride)
-		default:
-			return fmt.Errorf("unsupported API schema: backend=%s", out)
-		}
-		f.metrics = f.env.MessagesMetricsFactory.NewMetrics()
-	default:
-		return fmt.Errorf("unsupported endpoint for per-route upstream filter: %v", f.rf.endpoint)
-	}
-	return nil
-}