From 3af87de290228087159bf8653f0f8f96581af1f3 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 09:24:05 -0800 Subject: [PATCH 1/9] feat: implements CEL expression API for costs Signed-off-by: Takeshi Yoneda --- Makefile | 2 +- api/v1alpha1/api.go | 27 +++++-- go.mod | 3 + go.sum | 12 +++ internal/controller/sink.go | 8 ++ internal/controller/sink_test.go | 11 +++ internal/extproc/processor.go | 31 ++++++-- internal/extproc/processor_test.go | 26 +++++- internal/extproc/server.go | 19 ++++- internal/extproc/server_test.go | 24 ++++-- internal/llmcostcel/cel.go | 79 +++++++++++++++++++ internal/llmcostcel/cel_test.go | 60 ++++++++++++++ ...gateway.envoyproxy.io_aigatewayroutes.yaml | 23 ++++-- tests/cel-validation/main_test.go | 1 + .../testdata/aigatewayroutes/llmcosts.yaml | 33 ++++++++ tests/extproc/envoy.yaml | 1 + tests/extproc/extproc_test.go | 30 +++++-- 17 files changed, 354 insertions(+), 36 deletions(-) create mode 100644 internal/llmcostcel/cel.go create mode 100644 internal/llmcostcel/cel_test.go create mode 100644 tests/cel-validation/testdata/aigatewayroutes/llmcosts.yaml diff --git a/Makefile b/Makefile index 43ef37d81..0426cf243 100644 --- a/Makefile +++ b/Makefile @@ -103,7 +103,7 @@ editorconfig: editorconfig-checker .PHONY: test test: @echo "test => ./..." - @go test -v ./... + @go test -v -race ./... ENVTEST_K8S_VERSIONS ?= 1.29.0 1.30.0 1.31.0 diff --git a/api/v1alpha1/api.go b/api/v1alpha1/api.go index 3d7e598be..729c99b4f 100644 --- a/api/v1alpha1/api.go +++ b/api/v1alpha1/api.go @@ -475,17 +475,32 @@ type LLMRequestCost struct { // +kubebuilder:validation:Required MetadataKey string `json:"metadataKey"` // Type specifies the type of the request cost. The default is "OutputToken", - // and it uses "output token" as the cost. The other types are "InputToken" and "TotalToken". + // and it uses "output token" as the cost. The other types are "InputToken", "TotalToken", + // and "CEL". // - // +kubebuilder:validation:Enum=OutputToken;InputToken;TotalToken + // +kubebuilder:validation:Enum=OutputToken;InputToken;TotalToken;CEL Type LLMRequestCostType `json:"type"` // CELExpression is the CEL expression to calculate the cost of the request. - // The CEL expression must return an integer value. The CEL expression should be - // able to access the request headers, model name, backend name, input/output tokens etc. + // The CEL expression must return a signed or unsigned integer. If the + // return value is negative, it will be error. + // + // The expression can use the following variables: + // + // * model: the model name extracted from the request content. Type: string. + // * backend: the backend name in the form of "name.namespace". Type: string. + // * input_tokens: the number of input tokens. Type: unsigned integer. + // * output_tokens: the number of output tokens. Type: unsigned integer. + // * total_tokens: the total number of tokens. Type: unsigned integer. + // + // For example, the following expressions are valid: + // + // * "model == 'llama' ? input_tokens + output_token * 0.5 : total_tokens" + // * "backend == 'foo.default' ? input_tokens + output_tokens : total_tokens" + // * "input_tokens + output_tokens + total_tokens" + // * "input_tokens * output_tokens" // // +optional - // +notImplementedHide https://github.com/envoyproxy/ai-gateway/issues/97 - CELExpression *string `json:"celExpression"` + CELExpression *string `json:"celExpression,omitempty"` } // LLMRequestCostType specifies the type of the LLMRequestCost. diff --git a/go.mod b/go.mod index 34b8cbd2b..2f531805b 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/envoyproxy/gateway v0.5.0-rc.1.0.20250115172926-8b89dadfbd2c github.com/envoyproxy/go-control-plane/envoy v1.32.3 github.com/go-logr/logr v1.4.2 + github.com/google/cel-go v0.22.1 github.com/google/go-cmp v0.6.0 github.com/openai/openai-go v0.1.0-alpha.46 github.com/stretchr/testify v1.10.0 @@ -30,6 +31,7 @@ require ( require ( cel.dev/expr v0.18.0 // indirect + github.com/antlr4-go/antlr/v4 v4.13.0 // indirect github.com/aws/aws-sdk-go-v2/credentials v1.17.54 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.24 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.28 // indirect @@ -75,6 +77,7 @@ require ( github.com/prometheus/procfs v0.15.1 // indirect github.com/rogpeppe/go-internal v1.13.1 // indirect github.com/spf13/pflag v1.0.5 // indirect + github.com/stoewer/go-strcase v1.3.0 // indirect github.com/tidwall/gjson v1.14.4 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect diff --git a/go.sum b/go.sum index ffbb817df..d1d546530 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ cel.dev/expr v0.18.0 h1:CJ6drgk+Hf96lkLikr4rFf19WrU0BOWEihyZnI2TAzo= cel.dev/expr v0.18.0/go.mod h1:MrpN08Q+lEBs+bGYdLxxHkZoUSsCp0nSKTs0nTymJgw= +github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI= +github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g= github.com/aws/aws-sdk-go-v2 v1.33.0 h1:Evgm4DI9imD81V0WwD+TN4DCwjUMdc94TrduMLbgZJs= github.com/aws/aws-sdk-go-v2 v1.33.0/go.mod h1:P5WJBrYqqbWVaOxgH0X/FYYD47/nooaPOZPlQdmiN2U= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 h1:lL7IfaFzngfx0ZwUGOZdsFFnQ5uLvR0hWqqhyE7Q9M8= @@ -70,6 +72,8 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/cel-go v0.22.1 h1:AfVXx3chM2qwoSbM7Da8g8hX8OVSkBFwX+rz2+PcK40= +github.com/google/cel-go v0.22.1/go.mod h1:BuznPXXfQDpXKWQ9sPW3TzlAJN5zzFe+i9tIs0yC4s8= github.com/google/gnostic-models v0.6.9 h1:MU/8wDLif2qCXZmzncUQ/BOfxWfthHi63KqpoNbWqVw= github.com/google/gnostic-models v0.6.9/go.mod h1:CiWsm0s6BSQd1hRn8/QmxqB6BesYcbSZxsz9b0KuDBw= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= @@ -130,8 +134,15 @@ github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/stoewer/go-strcase v1.3.0 h1:g0eASXYtp+yvN9fK8sH94oCIk0fau9uV1/ZdJ0AVEzs= +github.com/stoewer/go-strcase v1.3.0/go.mod h1:fAH5hQ5pehh+j3nZfvwdk2RgEgQjAoM8wodgtPmh1xo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= @@ -224,6 +235,7 @@ gopkg.in/evanphx/json-patch.v4 v4.12.0 h1:n6jtcsulIzXPJaxegRbvFNNrZDjbij7ny3gmSP gopkg.in/evanphx/json-patch.v4 v4.12.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/api v0.32.1 h1:f562zw9cy+GvXzXf0CKlVQ7yHJVYzLfL6JAS4kOAaOc= diff --git a/internal/controller/sink.go b/internal/controller/sink.go index 7fdd88cd3..41de9dedb 100644 --- a/internal/controller/sink.go +++ b/internal/controller/sink.go @@ -14,6 +14,7 @@ import ( aigv1a1 "github.com/envoyproxy/ai-gateway/api/v1alpha1" "github.com/envoyproxy/ai-gateway/filterconfig" + "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) const selectedBackendHeaderKey = "x-ai-eg-selected-backend" @@ -203,6 +204,13 @@ func (c *configSink) updateExtProcConfigMap(aiGatewayRoute *aigv1a1.AIGatewayRou fc.Type = filterconfig.LLMRequestCostTypeTotalToken case aigv1a1.LLMRequestCostTypeCEL: fc.Type = filterconfig.LLMRequestCostTypeCELExpression + expr := *cost.CELExpression + // Sanity check the CEL expression. + _, err := llmcostcel.NewProgram(expr) + if err != nil { + return fmt.Errorf("invalid CEL expression: %w", err) + } + fc.CELExpression = expr default: return fmt.Errorf("unknown request cost type: %s", cost.Type) } diff --git a/internal/controller/sink_test.go b/internal/controller/sink_test.go index 792648915..5e461e04b 100644 --- a/internal/controller/sink_test.go +++ b/internal/controller/sink_test.go @@ -276,6 +276,15 @@ func Test_updateExtProcConfigMap(t *testing.T) { Type: aigv1a1.LLMRequestCostTypeInputToken, MetadataKey: "input-token", }, + { + Type: aigv1a1.LLMRequestCostTypeTotalToken, + MetadataKey: "total-token", + }, + { + Type: aigv1a1.LLMRequestCostTypeCEL, + MetadataKey: "cel-token", + CELExpression: ptr.To("model == 'cool_model' ? input_tokens * output_tokens : total_tokens"), + }, }, }, }, @@ -299,6 +308,8 @@ func Test_updateExtProcConfigMap(t *testing.T) { LLMRequestCosts: []filterconfig.LLMRequestCost{ {Type: filterconfig.LLMRequestCostTypeOutputToken, MetadataKey: "output-token"}, {Type: filterconfig.LLMRequestCostTypeInputToken, MetadataKey: "input-token"}, + {Type: filterconfig.LLMRequestCostTypeTotalToken, MetadataKey: "total-token"}, + {Type: filterconfig.LLMRequestCostTypeCELExpression, MetadataKey: "cel-token", CELExpression: "model == 'cool_model' ? input_tokens * output_tokens : total_tokens"}, }, }, }, diff --git a/internal/extproc/processor.go b/internal/extproc/processor.go index 410742b7f..8725a2aa6 100644 --- a/internal/extproc/processor.go +++ b/internal/extproc/processor.go @@ -11,6 +11,7 @@ import ( corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" + "github.com/google/cel-go/cel" "google.golang.org/protobuf/types/known/structpb" "github.com/envoyproxy/ai-gateway/extprocapi" @@ -18,6 +19,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/extproc/backendauth" "github.com/envoyproxy/ai-gateway/internal/extproc/router" "github.com/envoyproxy/ai-gateway/internal/extproc/translator" + "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) // processorConfig is the configuration for the processor. @@ -25,11 +27,16 @@ import ( type processorConfig struct { bodyParser router.RequestBodyParser router extprocapi.Router - ModelNameHeaderKey, selectedBackendHeaderKey string + modelNameHeaderKey, selectedBackendHeaderKey string factories map[filterconfig.VersionedAPISchema]translator.Factory backendAuthHandlers map[string]backendauth.Handler metadataNamespace string - requestCosts []filterconfig.LLMRequestCost + requestCosts []processorConfigRequestCost +} + +type processorConfigRequestCost struct { + *filterconfig.LLMRequestCost + celProg cel.Program } // ProcessorIface is the interface for the processor. @@ -79,7 +86,7 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt } p.logger.Info("Processing request", "path", path, "model", model) - p.requestHeaders[p.config.ModelNameHeaderKey] = model + p.requestHeaders[p.config.modelNameHeaderKey] = model b, err := p.config.router.Calculate(p.requestHeaders) if err != nil { return nil, fmt.Errorf("failed to calculate route: %w", err) @@ -107,7 +114,7 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt } // Set the model name to the request header with the key `x-ai-gateway-llm-model-name`. headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{ - Header: &corev3.HeaderValue{Key: p.config.ModelNameHeaderKey, RawValue: []byte(model)}, + Header: &corev3.HeaderValue{Key: p.config.modelNameHeaderKey, RawValue: []byte(model)}, }, &corev3.HeaderValueOption{ Header: &corev3.HeaderValue{Key: p.config.selectedBackendHeaderKey, RawValue: []byte(b.Name)}, }) @@ -203,7 +210,8 @@ func (p *Processor) ProcessResponseBody(_ context.Context, body *extprocv3.HttpB func (p *Processor) maybeBuildDynamicMetadata() (*structpb.Struct, error) { metadata := make(map[string]*structpb.Value, len(p.config.requestCosts)) - for _, c := range p.config.requestCosts { + for i := range p.config.requestCosts { + c := &p.config.requestCosts[i] var cost uint32 switch c.Type { case filterconfig.LLMRequestCostTypeInputToken: @@ -212,6 +220,19 @@ func (p *Processor) maybeBuildDynamicMetadata() (*structpb.Struct, error) { cost = p.costs.OutputTokens case filterconfig.LLMRequestCostTypeTotalToken: cost = p.costs.TotalTokens + case filterconfig.LLMRequestCostTypeCELExpression: + costU64, err := llmcostcel.EvaluateProgram( + c.celProg, + p.requestHeaders[p.config.modelNameHeaderKey], + p.requestHeaders[p.config.selectedBackendHeaderKey], + p.costs.InputTokens, + p.costs.OutputTokens, + p.costs.TotalTokens, + ) + if err != nil { + return nil, fmt.Errorf("failed to evaluate CEL expression: %w", err) + } + cost = uint32(costU64) //nolint:gosec default: return nil, fmt.Errorf("unknown request cost kind: %s", c.Type) } diff --git a/internal/extproc/processor_test.go b/internal/extproc/processor_test.go index 92f834875..6a37c2fc3 100644 --- a/internal/extproc/processor_test.go +++ b/internal/extproc/processor_test.go @@ -13,6 +13,7 @@ import ( "github.com/envoyproxy/ai-gateway/filterconfig" "github.com/envoyproxy/ai-gateway/internal/extproc/router" "github.com/envoyproxy/ai-gateway/internal/extproc/translator" + "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) func TestProcessor_ProcessRequestHeaders(t *testing.T) { @@ -65,11 +66,24 @@ func TestProcessor_ProcessResponseBody(t *testing.T) { retBodyMutation: expBodyMut, retHeaderMutation: expHeadMut, retUsedToken: translator.LLMTokenUsage{OutputTokens: 123, InputTokens: 1}, } + + celProgInt, err := llmcostcel.NewProgram("54321") + require.NoError(t, err) + celProgUint, err := llmcostcel.NewProgram("uint(9999)") + require.NoError(t, err) p := &Processor{translator: mt, config: &processorConfig{ metadataNamespace: "ai_gateway_llm_ns", - requestCosts: []filterconfig.LLMRequestCost{ - {Type: filterconfig.LLMRequestCostTypeOutputToken, MetadataKey: "output_token_usage"}, - {Type: filterconfig.LLMRequestCostTypeInputToken, MetadataKey: "input_token_usage"}, + requestCosts: []processorConfigRequestCost{ + {LLMRequestCost: &filterconfig.LLMRequestCost{Type: filterconfig.LLMRequestCostTypeOutputToken, MetadataKey: "output_token_usage"}}, + {LLMRequestCost: &filterconfig.LLMRequestCost{Type: filterconfig.LLMRequestCostTypeInputToken, MetadataKey: "input_token_usage"}}, + { + celProg: celProgInt, + LLMRequestCost: &filterconfig.LLMRequestCost{Type: filterconfig.LLMRequestCostTypeCELExpression, MetadataKey: "cel_int"}, + }, + { + celProg: celProgUint, + LLMRequestCost: &filterconfig.LLMRequestCost{Type: filterconfig.LLMRequestCostTypeCELExpression, MetadataKey: "cel_uint"}, + }, }, }} res, err := p.ProcessResponseBody(context.Background(), inBody) @@ -84,6 +98,10 @@ func TestProcessor_ProcessResponseBody(t *testing.T) { GetStructValue().Fields["output_token_usage"].GetNumberValue()) require.Equal(t, float64(1), md.Fields["ai_gateway_llm_ns"]. GetStructValue().Fields["input_token_usage"].GetNumberValue()) + require.Equal(t, float64(54321), md.Fields["ai_gateway_llm_ns"]. + GetStructValue().Fields["cel_int"].GetNumberValue()) + require.Equal(t, float64(9999), md.Fields["ai_gateway_llm_ns"]. + GetStructValue().Fields["cel_uint"].GetNumberValue()) }) } @@ -168,7 +186,7 @@ func TestProcessor_ProcessRequestBody(t *testing.T) { {Name: "some-schema", Version: "v10.0"}: factory.impl, }, selectedBackendHeaderKey: "x-ai-gateway-backend-key", - ModelNameHeaderKey: "x-ai-gateway-model-key", + modelNameHeaderKey: "x-ai-gateway-model-key", }, requestHeaders: headers, logger: slog.Default()} resp, err := p.ProcessRequestBody(context.Background(), &extprocv3.HttpBody{}) require.NoError(t, err) diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 2f647c072..7ffce36e2 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -8,6 +8,7 @@ import ( "log/slog" extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" + "github.com/google/cel-go/cel" "google.golang.org/grpc/codes" "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/status" @@ -17,6 +18,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/extproc/backendauth" "github.com/envoyproxy/ai-gateway/internal/extproc/router" "github.com/envoyproxy/ai-gateway/internal/extproc/translator" + "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) // Server implements the external process server. @@ -64,14 +66,27 @@ func (s *Server[P]) LoadConfig(config *filterconfig.Config) error { } } + costs := make([]processorConfigRequestCost, 0, len(config.LLMRequestCosts)) + for i := range config.LLMRequestCosts { + c := &config.LLMRequestCosts[i] + var prog cel.Program + if c.CELExpression != "" { + prog, err = llmcostcel.NewProgram(c.CELExpression) + if err != nil { + return fmt.Errorf("cannot create CEL program for cost: %w", err) + } + } + costs = append(costs, processorConfigRequestCost{LLMRequestCost: c, celProg: prog}) + } + newConfig := &processorConfig{ bodyParser: bodyParser, router: rt, selectedBackendHeaderKey: config.SelectedBackendHeaderKey, - ModelNameHeaderKey: config.ModelNameHeaderKey, + modelNameHeaderKey: config.ModelNameHeaderKey, factories: factories, backendAuthHandlers: backendAuthHandlers, metadataNamespace: config.MetadataNamespace, - requestCosts: config.LLMRequestCosts, + requestCosts: costs, } s.config = newConfig // This is racey, but we don't care. return nil diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index 2c4f06816..fd8c41e1d 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -16,6 +16,7 @@ import ( "google.golang.org/grpc/status" "github.com/envoyproxy/ai-gateway/filterconfig" + "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) func requireNewServerWithMockProcessor(t *testing.T) *Server[*mockProcessor] { @@ -36,8 +37,11 @@ func TestServer_LoadConfig(t *testing.T) { }) t.Run("ok", func(t *testing.T) { config := &filterconfig.Config{ - MetadataNamespace: "ns", - LLMRequestCosts: []filterconfig.LLMRequestCost{{MetadataKey: "key", Type: filterconfig.LLMRequestCostTypeOutputToken}}, + MetadataNamespace: "ns", + LLMRequestCosts: []filterconfig.LLMRequestCost{ + {MetadataKey: "key", Type: filterconfig.LLMRequestCostTypeOutputToken}, + {MetadataKey: "cel_key", Type: filterconfig.LLMRequestCostTypeCELExpression, CELExpression: "1 + 1"}, + }, Schema: filterconfig.VersionedAPISchema{Name: filterconfig.APISchemaOpenAI}, SelectedBackendHeaderKey: "x-ai-eg-selected-backend", ModelNameHeaderKey: "x-model-name", @@ -72,17 +76,25 @@ func TestServer_LoadConfig(t *testing.T) { require.NoError(t, err) require.NotNil(t, s.config) - require.NotEmpty(t, s.config.requestCosts) require.Equal(t, "ns", s.config.metadataNamespace) - require.Equal(t, "key", s.config.requestCosts[0].MetadataKey) - require.Equal(t, filterconfig.LLMRequestCostTypeOutputToken, s.config.requestCosts[0].Type) require.NotNil(t, s.config.router) require.NotNil(t, s.config.bodyParser) require.Equal(t, "x-ai-eg-selected-backend", s.config.selectedBackendHeaderKey) - require.Equal(t, "x-model-name", s.config.ModelNameHeaderKey) + require.Equal(t, "x-model-name", s.config.modelNameHeaderKey) require.Len(t, s.config.factories, 2) require.NotNil(t, s.config.factories[filterconfig.VersionedAPISchema{Name: filterconfig.APISchemaOpenAI}]) require.NotNil(t, s.config.factories[filterconfig.VersionedAPISchema{Name: filterconfig.APISchemaAWSBedrock}]) + + require.Len(t, s.config.requestCosts, 2) + require.Equal(t, filterconfig.LLMRequestCostTypeOutputToken, s.config.requestCosts[0].Type) + require.Equal(t, "key", s.config.requestCosts[0].MetadataKey) + require.Equal(t, filterconfig.LLMRequestCostTypeCELExpression, s.config.requestCosts[1].Type) + require.Equal(t, "1 + 1", s.config.requestCosts[1].CELExpression) + prog := s.config.requestCosts[1].celProg + require.NotNil(t, prog) + val, err := llmcostcel.EvaluateProgram(prog, "", "", 1, 1, 1) + require.NoError(t, err) + require.Equal(t, uint64(2), val) }) } diff --git a/internal/llmcostcel/cel.go b/internal/llmcostcel/cel.go new file mode 100644 index 000000000..77e1015e2 --- /dev/null +++ b/internal/llmcostcel/cel.go @@ -0,0 +1,79 @@ +// Package llmcostcel provides functions to create and evaluate CEL programs to calculate costs. +// +// This exists as a separate package to be used both in the controller to validate the expression +// and in the external processor to evaluate the expression. +package llmcostcel + +import ( + "fmt" + + "github.com/google/cel-go/cel" +) + +const ( + celModelNameKey = "model" + celBackendKey = "backend" + celInputTokensKey = "input_tokens" + celOutputTokensKey = "output_tokens" + celTotalTokensKey = "total_tokens" +) + +// NewProgram creates a new CEL program from the given expression. +func NewProgram(expr string) (prog cel.Program, err error) { + env, err := cel.NewEnv( + cel.Variable(celModelNameKey, cel.StringType), + cel.Variable(celBackendKey, cel.StringType), + cel.Variable(celInputTokensKey, cel.UintType), + cel.Variable(celOutputTokensKey, cel.UintType), + cel.Variable(celTotalTokensKey, cel.UintType), + ) + if err != nil { + return nil, fmt.Errorf("cannot create CEL environment: %w", err) + } + ast, issues := env.Compile(expr) + if issues != nil && issues.Err() != nil { + err := issues.Err() + return nil, fmt.Errorf("cannot compile CEL expression: %w", err) + } + prog, err = env.Program(ast) + if err != nil { + return nil, fmt.Errorf("cannot create CEL program: %w", err) + } + + // Sanity check by evaluating the expression with some dummy values. + _, err = EvaluateProgram(prog, "dummy", "dummy", 0, 0, 0) + if err != nil { + return nil, fmt.Errorf("failed to evaluate CEL expression: %w", err) + } + return prog, nil +} + +// EvaluateProgram evaluates the given CEL program with the given variables. +func EvaluateProgram(prog cel.Program, modelName, backend string, inputTokens, outputTokens, totalTokens uint32) (uint64, error) { + out, _, err := prog.Eval(map[string]interface{}{ + celModelNameKey: modelName, + celBackendKey: backend, + celInputTokensKey: inputTokens, + celOutputTokensKey: outputTokens, + celTotalTokensKey: totalTokens, + }) + if err != nil { + return 0, fmt.Errorf("failed to evaluate CEL expression: %w", err) + } + if out == nil { + return 0, fmt.Errorf("CEL expression result is nil") + } + + switch out.Type() { + case cel.IntType: + result := out.Value().(int64) + if result < 0 { + return 0, fmt.Errorf("CEL expression result is negative") + } + return uint64(result), nil + case cel.UintType: + return out.Value().(uint64), nil + default: + return 0, fmt.Errorf("CEL expression result is not an integer") + } +} diff --git a/internal/llmcostcel/cel_test.go b/internal/llmcostcel/cel_test.go new file mode 100644 index 000000000..2449e8f22 --- /dev/null +++ b/internal/llmcostcel/cel_test.go @@ -0,0 +1,60 @@ +package llmcostcel + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func Test_newCelProgram(t *testing.T) { + t.Run("int", func(t *testing.T) { + _, err := NewProgram("1 + 1") + require.NoError(t, err) + }) + t.Run("uint", func(t *testing.T) { + _, err := NewProgram("uint(1) + uint(1)") + require.NoError(t, err) + }) + t.Run("use cool_model", func(t *testing.T) { + prog, err := NewProgram("model == 'cool_model' ? input_tokens * output_tokens : total_tokens") + require.NoError(t, err) + out, _, err := prog.Eval(map[string]interface{}{ + celModelNameKey: "cool_model", + celBackendKey: "cool_backend", + celInputTokensKey: uint(100), + celOutputTokensKey: uint(2), + celTotalTokensKey: uint(3), + }) + require.NoError(t, err) + require.Equal(t, uint64(200), out.Value().(uint64)) + + out, _, err = prog.Eval(map[string]interface{}{ + celModelNameKey: "not_cool_model", + celBackendKey: "cool_backend", + celInputTokensKey: uint(1), + celOutputTokensKey: uint(2), + celTotalTokensKey: uint(3), + }) + require.NoError(t, err) + require.Equal(t, uint64(3), out.Value().(uint64)) + }) + + t.Run("ensure concurrency safety", func(t *testing.T) { + prog, err := NewProgram("model == 'cool_model' ? input_tokens * output_tokens : total_tokens") + require.NoError(t, err) + + // Ensure that the program can be evaluated concurrently. + var wg sync.WaitGroup + wg.Add(100) + for i := 0; i < 100; i++ { + go func() { + defer wg.Done() + v, err := EvaluateProgram(prog, "cool_model", "cool_backend", 100, 2, 3) + require.NoError(t, err) + require.Equal(t, uint64(200), v) + }() + } + wg.Wait() + }) +} diff --git a/manifests/charts/ai-gateway-helm/crds/aigateway.envoyproxy.io_aigatewayroutes.yaml b/manifests/charts/ai-gateway-helm/crds/aigateway.envoyproxy.io_aigatewayroutes.yaml index 038df7dd3..122fae5e2 100644 --- a/manifests/charts/ai-gateway-helm/crds/aigateway.envoyproxy.io_aigatewayroutes.yaml +++ b/manifests/charts/ai-gateway-helm/crds/aigateway.envoyproxy.io_aigatewayroutes.yaml @@ -196,10 +196,21 @@ spec: description: LLMRequestCost configures each request cost. properties: celExpression: - description: |- - CELExpression is the CEL expression to calculate the cost of the request. - The CEL expression must return an integer value. The CEL expression should be - able to access the request headers, model name, backend name, input/output tokens etc. + description: "CELExpression is the CEL expression to calculate + the cost of the request.\nThe CEL expression must return a + signed or unsigned integer. If the\nreturn value is negative, + it will be error.\n\nThe expression can use the following + variables:\n\n\t* model: the model name extracted from the + request content. Type: string.\n\t* backend: the backend name + in the form of \"name.namespace\". Type: string.\n\t* input_tokens: + the number of input tokens. Type: unsigned integer.\n\t* output_tokens: + the number of output tokens. Type: unsigned integer.\n\t* + total_tokens: the total number of tokens. Type: unsigned integer.\n\nFor + example, the following expressions are valid:\n\n\t* \"model + == 'llama' ? input_tokens + output_token * 0.5 : total_tokens\"\n\t* + \"backend == 'foo.default' ? input_tokens + output_tokens + : total_tokens\"\n\t* \"input_tokens + output_tokens + total_tokens\"\n\t* + \"input_tokens * output_tokens\"" type: string metadataKey: description: MetadataKey is the key of the metadata to store @@ -208,11 +219,13 @@ spec: type: description: |- Type specifies the type of the request cost. The default is "OutputToken", - and it uses "output token" as the cost. The other types are "InputToken" and "TotalToken". + and it uses "output token" as the cost. The other types are "InputToken", "TotalToken", + and "CEL". enum: - OutputToken - InputToken - TotalToken + - CEL type: string required: - metadataKey diff --git a/tests/cel-validation/main_test.go b/tests/cel-validation/main_test.go index 5f5da4a0a..8452f1b2e 100644 --- a/tests/cel-validation/main_test.go +++ b/tests/cel-validation/main_test.go @@ -29,6 +29,7 @@ func TestAIGatewayRoutes(t *testing.T) { expErr string }{ {name: "basic.yaml"}, + {name: "llmcosts.yaml"}, { name: "non_openai_schema.yaml", expErr: `spec.schema: Invalid value: "object": failed rule: self.name == 'OpenAI'`, diff --git a/tests/cel-validation/testdata/aigatewayroutes/llmcosts.yaml b/tests/cel-validation/testdata/aigatewayroutes/llmcosts.yaml new file mode 100644 index 000000000..07160b2c2 --- /dev/null +++ b/tests/cel-validation/testdata/aigatewayroutes/llmcosts.yaml @@ -0,0 +1,33 @@ +apiVersion: aigateway.envoyproxy.io/v1alpha1 +kind: AIGatewayRoute +metadata: + name: llmcosts + namespace: default +spec: + schema: + name: OpenAI + targetRefs: + - name: some-gateway + kind: Gateway + group: gateway.networking.k8s.io + rules: + - matches: + - headers: + - type: Exact + name: x-ai-eg-model + value: llama3-70b + backendRefs: + - name: kserve + weight: 20 + - name: aws-bedrock + weight: 80 + llmRequestCosts: + - metadataKey: llm_input_token + type: InputToken + - metadataKey: llm_output_token + type: OutputToken + - metadataKey: llm_total_token + type: TotalToken + - metadataKey: some_cel_cost + type: CEL + celExpression: "llm_input_token + llm_output_token + llm_total_token" diff --git a/tests/extproc/envoy.yaml b/tests/extproc/envoy.yaml index ea7c90510..059d8833f 100644 --- a/tests/extproc/envoy.yaml +++ b/tests/extproc/envoy.yaml @@ -18,6 +18,7 @@ static_resources: path: ACCESS_LOG_PATH json_format: "used_token": "%DYNAMIC_METADATA(ai_gateway_llm_ns:used_token)%" + "some_cel": "%DYNAMIC_METADATA(ai_gateway_llm_ns:some_cel)%" route_config: virtual_hosts: - name: local_route diff --git a/tests/extproc/extproc_test.go b/tests/extproc/extproc_test.go index 72fb61015..2af714dda 100644 --- a/tests/extproc/extproc_test.go +++ b/tests/extproc/extproc_test.go @@ -54,6 +54,7 @@ func TestE2E(t *testing.T) { MetadataNamespace: "ai_gateway_llm_ns", LLMRequestCosts: []filterconfig.LLMRequestCost{ {MetadataKey: "used_token", Type: filterconfig.LLMRequestCostTypeInputToken}, + {MetadataKey: "some_cel", Type: filterconfig.LLMRequestCostTypeCELExpression, CELExpression: "1 + 1"}, }, Schema: openAISchema, // This can be any header key, but it must match the envoy.yaml routing configuration. @@ -118,6 +119,7 @@ func TestE2E(t *testing.T) { // This should match the format of the access log in envoy.yaml. type lineFormat struct { UsedToken any `json:"used_token"` + SomeCel any `json:"some_cel"` } scanner := bufio.NewScanner(bytes.NewReader(accessLog)) for scanner.Scan() { @@ -130,14 +132,15 @@ func TestE2E(t *testing.T) { t.Logf("line: %s", line) // The access formatter somehow changed its behavior sometimes between 1.31 and the latest Envoy, // so we need to check for both float64 and string. - if num, ok := l.UsedToken.(float64); ok && num > 0 { - return true - } else if str, ok := l.UsedToken.(string); ok { - if num, err := strconv.Atoi(str); err == nil && num > 0 { - return true - } + if !anyCostGreaterThanZero(l.SomeCel) { + t.Log("some_cel is not existent or greater than zero") + continue + } + if !anyCostGreaterThanZero(l.UsedToken) { + t.Log("used_token is not existent or greater than zero") + continue } - t.Log("cannot find used token in line") + return true } return false }, 10*time.Second, 1*time.Second) @@ -274,6 +277,19 @@ func requireBinaries(t *testing.T) { } } +func anyCostGreaterThanZero(cost any) bool { + // The access formatter somehow changed its behavior sometimes between 1.31 and the latest Envoy, + // so we need to check for both float64 and string. + if num, ok := cost.(float64); ok && num > 0 { + return true + } else if str, ok := cost.(string); ok { + if num, err := strconv.Atoi(str); err == nil && num > 0 { + return true + } + } + return false +} + // getEnvVarOrSkip requires an environment variable to be set. func getEnvVarOrSkip(t *testing.T, envVar string) string { value := os.Getenv(envVar) From db900a405e13cdeb971fb7a86e063c7e704e7ded Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 09:29:31 -0800 Subject: [PATCH 2/9] more Signed-off-by: Takeshi Yoneda --- internal/llmcostcel/cel.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/llmcostcel/cel.go b/internal/llmcostcel/cel.go index 77e1015e2..8c7331343 100644 --- a/internal/llmcostcel/cel.go +++ b/internal/llmcostcel/cel.go @@ -68,12 +68,12 @@ func EvaluateProgram(prog cel.Program, modelName, backend string, inputTokens, o case cel.IntType: result := out.Value().(int64) if result < 0 { - return 0, fmt.Errorf("CEL expression result is negative") + return 0, fmt.Errorf("CEL expression result is negative (%d)", result) } return uint64(result), nil case cel.UintType: return out.Value().(uint64), nil default: - return 0, fmt.Errorf("CEL expression result is not an integer") + return 0, fmt.Errorf("CEL expression result is not an integer, got %v", out.Type()) } } From 4567b1a9cf66e98cba2366d3b6a3cee9a2eeda83 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 09:31:44 -0800 Subject: [PATCH 3/9] more Signed-off-by: Takeshi Yoneda --- internal/llmcostcel/cel_test.go | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/internal/llmcostcel/cel_test.go b/internal/llmcostcel/cel_test.go index 2449e8f22..1e30da1b8 100644 --- a/internal/llmcostcel/cel_test.go +++ b/internal/llmcostcel/cel_test.go @@ -7,7 +7,7 @@ import ( "github.com/stretchr/testify/require" ) -func Test_newCelProgram(t *testing.T) { +func TestNewProgram(t *testing.T) { t.Run("int", func(t *testing.T) { _, err := NewProgram("1 + 1") require.NoError(t, err) @@ -16,28 +16,16 @@ func Test_newCelProgram(t *testing.T) { _, err := NewProgram("uint(1) + uint(1)") require.NoError(t, err) }) - t.Run("use cool_model", func(t *testing.T) { + t.Run("variables", func(t *testing.T) { prog, err := NewProgram("model == 'cool_model' ? input_tokens * output_tokens : total_tokens") require.NoError(t, err) - out, _, err := prog.Eval(map[string]interface{}{ - celModelNameKey: "cool_model", - celBackendKey: "cool_backend", - celInputTokensKey: uint(100), - celOutputTokensKey: uint(2), - celTotalTokensKey: uint(3), - }) + v, err := EvaluateProgram(prog, "cool_model", "cool_backend", 100, 2, 3) require.NoError(t, err) - require.Equal(t, uint64(200), out.Value().(uint64)) + require.Equal(t, uint64(200), v) - out, _, err = prog.Eval(map[string]interface{}{ - celModelNameKey: "not_cool_model", - celBackendKey: "cool_backend", - celInputTokensKey: uint(1), - celOutputTokensKey: uint(2), - celTotalTokensKey: uint(3), - }) + v, err = EvaluateProgram(prog, "not_cool_model", "cool_backend", 100, 2, 3) require.NoError(t, err) - require.Equal(t, uint64(3), out.Value().(uint64)) + require.Equal(t, uint64(3), v) }) t.Run("ensure concurrency safety", func(t *testing.T) { From 569ac5d1dd22d19ef05ebe84c6cdef8b2726b8c2 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 14:31:44 -0800 Subject: [PATCH 4/9] more Signed-off-by: Takeshi Yoneda --- Makefile | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 0426cf243..589618bb9 100644 --- a/Makefile +++ b/Makefile @@ -19,8 +19,10 @@ ENABLE_MULTI_PLATFORMS ?= false HELM_CHART_VERSION ?= v0.0.0-latest # Arguments for go test. This can be used, for example, to run specific tests via -# `GO_TEST_EXTRA_ARGS="-run TestName/foo/etc"`. -GO_TEST_EXTRA_ARGS ?= +# `GO_TEST_ARGS="-run TestName/foo/etc" -v -race`. +GO_TEST_ARGS ?= -v -race +# Arguments for go test in e2e tests in addition to GO_TEST_ARGS, applicable to test-e2e, test-extproc, and test-controller. +GO_TEST_E2E_ARGS ?= -count=1 # This will print out the help message for contributing to the project. .PHONY: help @@ -115,7 +117,7 @@ test-cel: envtest apigen @for k8sVersion in $(ENVTEST_K8S_VERSIONS); do \ echo "Run CEL Validation on k8s $$k8sVersion"; \ KUBEBUILDER_ASSETS="$$($(ENVTEST) use $$k8sVersion -p path)" \ - go test ./tests/cel-validation $(GO_TEST_EXTRA_ARGS) --tags test_cel_validation -v -count=1; \ + go test ./tests/cel-validation $(GO_TEST_ARGS) $(GO_TEST_E2E_ARGS) --tags test_cel_validation; \ done # This runs the end-to-end tests for extproc without controller or k8s at all. @@ -127,7 +129,7 @@ test-extproc: build.extproc @$(MAKE) build.extproc_custom_router CMD_PATH_PREFIX=examples @$(MAKE) build.testupstream CMD_PATH_PREFIX=tests @echo "Run ExtProc test" - @go test ./tests/extproc/... $(GO_TEST_EXTRA_ARGS) -tags test_extproc -v -count=1 + @go test ./tests/extproc/... $(GO_TEST_ARGS) $(GO_TEST_E2E_ARGS) -tags test_extproc -v # This runs the end-to-end tests for the controller with EnvTest. .PHONY: test-controller @@ -135,7 +137,7 @@ test-controller: envtest apigen @for k8sVersion in $(ENVTEST_K8S_VERSIONS); do \ echo "Run Controller tests on k8s $$k8sVersion"; \ KUBEBUILDER_ASSETS="$$($(ENVTEST) use $$k8sVersion -p path)" \ - go test ./tests/controller $(GO_TEST_EXTRA_ARGS) --tags test_controller -v -count=1; \ + go test ./tests/controller $(GO_TEST_ARGS) $(GO_TEST_E2E_ARGS) -tags test_controller; \ done # This runs the end-to-end tests for the controller and extproc with a local kind cluster. @@ -146,7 +148,7 @@ test-e2e: kind @$(MAKE) docker-build DOCKER_BUILD_ARGS="--load" @$(MAKE) docker-build.testupstream CMD_PATH_PREFIX=tests DOCKER_BUILD_ARGS="--load" @echo "Run E2E tests" - @go test ./tests/e2e/... $(GO_TEST_EXTRA_ARGS) -tags test_e2e -v -count=1 + @go test ./tests/e2e/... $(GO_TEST_ARGS) $(GO_TEST_E2E_ARGS) -tags test_e2e # This builds a binary for the given command under the internal/cmd directory. # From 5eac3891f285a4232fc21330276536a1350f9d31 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 14:32:12 -0800 Subject: [PATCH 5/9] more Signed-off-by: Takeshi Yoneda --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 589618bb9..503f2774b 100644 --- a/Makefile +++ b/Makefile @@ -105,7 +105,7 @@ editorconfig: editorconfig-checker .PHONY: test test: @echo "test => ./..." - @go test -v -race ./... + @go test $(GO_TEST_ARGS) ./... ENVTEST_K8S_VERSIONS ?= 1.29.0 1.30.0 1.31.0 From 6cb71ab11da7e5f7089360c5c4cb24d73cf64e45 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 14:44:49 -0800 Subject: [PATCH 6/9] more Signed-off-by: Takeshi Yoneda --- tests/extproc/custom_extproc_test.go | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/extproc/custom_extproc_test.go b/tests/extproc/custom_extproc_test.go index 0f397ef79..764121ee8 100644 --- a/tests/extproc/custom_extproc_test.go +++ b/tests/extproc/custom_extproc_test.go @@ -3,10 +3,10 @@ package extproc import ( - "bytes" "context" "encoding/base64" "fmt" + "os" "runtime" "testing" "time" @@ -36,8 +36,13 @@ func TestExtProcCustomRouter(t *testing.T) { }, }, }) - stdout := &bytes.Buffer{} - requireExtProc(t, stdout, fmt.Sprintf("../../out/extproc_custom_router-%s-%s", + stdoutPath := t.TempDir() + "/extproc-stdout.log" + f, err := os.Create(stdoutPath) + require.NoError(t, err) + defer func() { + require.NoError(t, f.Close()) + }() + requireExtProc(t, f, fmt.Sprintf("../../out/extproc_custom_router-%s-%s", runtime.GOOS, runtime.GOARCH), configPath) require.Eventually(t, func() bool { @@ -63,5 +68,12 @@ func TestExtProcCustomRouter(t *testing.T) { return true }, 10*time.Second, 1*time.Second) - require.Contains(t, stdout.String(), "model name: something-cool") // This must be logged by the custom router. + // Check that the custom router logs the model name after the file is closed. + defer func() { + + stdout, err := os.ReadFile(stdoutPath) + require.NoError(t, err) + t.Logf("stdout: %s", stdout) + require.Contains(t, string(stdout), "model name: something-cool") // This must be logged by the custom router. + }() } From 7777f5f9fbfa7a6ba3d35834dae339d6f969f210 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 14:45:19 -0800 Subject: [PATCH 7/9] more Signed-off-by: Takeshi Yoneda --- tests/extproc/custom_extproc_test.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/extproc/custom_extproc_test.go b/tests/extproc/custom_extproc_test.go index 764121ee8..98bb0b7d9 100644 --- a/tests/extproc/custom_extproc_test.go +++ b/tests/extproc/custom_extproc_test.go @@ -11,7 +11,7 @@ import ( "testing" "time" - openai "github.com/openai/openai-go" + "github.com/openai/openai-go" "github.com/openai/openai-go/option" "github.com/stretchr/testify/require" @@ -70,7 +70,6 @@ func TestExtProcCustomRouter(t *testing.T) { // Check that the custom router logs the model name after the file is closed. defer func() { - stdout, err := os.ReadFile(stdoutPath) require.NoError(t, err) t.Logf("stdout: %s", stdout) From ad6bff67f555f91f7a9fab2787e5f3e0975a730b Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Tue, 21 Jan 2025 15:09:38 -0800 Subject: [PATCH 8/9] more Signed-off-by: Takeshi Yoneda --- tests/extproc/real_providers_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/extproc/real_providers_test.go b/tests/extproc/real_providers_test.go index 7f102fd78..8e5673c5f 100644 --- a/tests/extproc/real_providers_test.go +++ b/tests/extproc/real_providers_test.go @@ -47,6 +47,7 @@ func TestWithRealProviders(t *testing.T) { MetadataNamespace: "ai_gateway_llm_ns", LLMRequestCosts: []filterconfig.LLMRequestCost{ {MetadataKey: "used_token", Type: filterconfig.LLMRequestCostTypeInputToken}, + {MetadataKey: "some_cel", Type: filterconfig.LLMRequestCostTypeCELExpression, CELExpression: "1+1"}, }, Schema: openAISchema, // This can be any header key, but it must match the envoy.yaml routing configuration. From 622c7cbadf7c555a93eeebc9783ccf18073e1e3b Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Thu, 23 Jan 2025 21:22:05 -0800 Subject: [PATCH 9/9] fix makefile comment Signed-off-by: Takeshi Yoneda --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 503f2774b..815ef9efa 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ ENABLE_MULTI_PLATFORMS ?= false HELM_CHART_VERSION ?= v0.0.0-latest # Arguments for go test. This can be used, for example, to run specific tests via -# `GO_TEST_ARGS="-run TestName/foo/etc" -v -race`. +# `GO_TEST_ARGS="-run TestName/foo/etc -v -race"`. GO_TEST_ARGS ?= -v -race # Arguments for go test in e2e tests in addition to GO_TEST_ARGS, applicable to test-e2e, test-extproc, and test-controller. GO_TEST_E2E_ARGS ?= -count=1