Skip to content

Commit 3a490bf

Browse files
mathetakemathetake
authored
extproc: adds ability to insert custom router (#104)
Signed-off-by: Takeshi Yoneda <[email protected]>
1 parent 08c31fd commit 3a490bf

File tree

15 files changed

+351
-120
lines changed

15 files changed

+351
-120
lines changed

‎Makefile

+2
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ test-cel: envtest apigen
113113
# This requires the extproc binary to be built as well as Envoy binary to be available in the PATH.
114114
.PHONY: test-extproc # This requires the extproc binary to be built.
115115
test-extproc: build.extproc
116+
@$(MAKE) build.extproc_custom_router CMD_PATH_PREFIX=examples
116117
@$(MAKE) build.testupstream CMD_PATH_PREFIX=tests
117118
@echo "Run ExtProc test"
118119
@go test ./tests/extproc/... -tags test_extproc -v -count=1
@@ -140,6 +141,7 @@ test-e2e: kind
140141
# Example:
141142
# - `make build.controller`: will build the cmd/controller directory.
142143
# - `make build.extproc`: will build the cmd/extproc directory.
144+
# - `make build.extproc_custom_router CMD_PATH_PREFIX=examples`: will build the examples/extproc_custom_router directory.
143145
# - `make build.testupstream CMD_PATH_PREFIX=tests`: will build the tests/testupstream directory.
144146
#
145147
# By default, this will build for the current GOOS and GOARCH.

‎cmd/extproc/main.go

+2-75
Original file line numberDiff line numberDiff line change
@@ -1,78 +1,5 @@
11
package main
22

3-
import (
4-
"context"
5-
"flag"
6-
"log"
7-
"log/slog"
8-
"net"
9-
"os"
10-
"os/signal"
11-
"syscall"
12-
"time"
3+
import "github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
134

14-
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
15-
"google.golang.org/grpc"
16-
"google.golang.org/grpc/health/grpc_health_v1"
17-
18-
"github.com/envoyproxy/ai-gateway/internal/extproc"
19-
"github.com/envoyproxy/ai-gateway/internal/version"
20-
)
21-
22-
var (
23-
configPath = flag.String("configPath", "", "path to the configuration file. "+
24-
"The file must be in YAML format specified in extprocconfig.Config type. The configuration file is watched for changes.")
25-
// TODO: unix domain socket support.
26-
extProcPort = flag.String("extProcPort", ":1063", "gRPC port for the external processor")
27-
logLevel = flag.String("logLevel", "info", "log level")
28-
)
29-
30-
func main() {
31-
flag.Parse()
32-
33-
var level slog.Level
34-
if err := level.UnmarshalText([]byte(*logLevel)); err != nil {
35-
log.Fatalf("failed to unmarshal log level: %v", err)
36-
}
37-
l := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
38-
Level: level,
39-
}))
40-
41-
l.Info("starting external processor", slog.String("version", version.Version))
42-
43-
if *configPath == "" {
44-
log.Fatal("configPath must be provided")
45-
}
46-
47-
ctx, cancel := context.WithCancel(context.Background())
48-
signalsChan := make(chan os.Signal, 1)
49-
signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
50-
go func() {
51-
<-signalsChan
52-
cancel()
53-
}()
54-
55-
// TODO: unix domain socket support.
56-
lis, err := net.Listen("tcp", *extProcPort)
57-
if err != nil {
58-
log.Fatalf("failed to listen: %v", err)
59-
}
60-
61-
server, err := extproc.NewServer[*extproc.Processor](l, extproc.NewProcessor)
62-
if err != nil {
63-
log.Fatalf("failed to create external processor server: %v", err)
64-
}
65-
66-
if err := extproc.StartConfigWatcher(ctx, *configPath, server, l, time.Second*5); err != nil {
67-
log.Fatalf("failed to start config watcher: %v", err)
68-
}
69-
70-
s := grpc.NewServer()
71-
extprocv3.RegisterExternalProcessorServer(s, server)
72-
grpc_health_v1.RegisterHealthServer(s, server)
73-
go func() {
74-
<-ctx.Done()
75-
s.GracefulStop()
76-
}()
77-
_ = s.Serve(lis)
78-
}
5+
func main() { mainlib.Main() }

‎cmd/extproc/mainlib/main.go

+80
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package mainlib
2+
3+
import (
4+
"context"
5+
"flag"
6+
"log"
7+
"log/slog"
8+
"net"
9+
"os"
10+
"os/signal"
11+
"syscall"
12+
"time"
13+
14+
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
15+
"google.golang.org/grpc"
16+
"google.golang.org/grpc/health/grpc_health_v1"
17+
18+
"github.com/envoyproxy/ai-gateway/internal/extproc"
19+
"github.com/envoyproxy/ai-gateway/internal/version"
20+
)
21+
22+
var (
23+
configPath = flag.String("configPath", "", "path to the configuration file. "+
24+
"The file must be in YAML format specified in extprocconfig.Config type. The configuration file is watched for changes.")
25+
// TODO: unix domain socket support.
26+
extProcPort = flag.String("extProcPort", ":1063", "gRPC port for the external processor")
27+
logLevel = flag.String("logLevel", "info", "log level")
28+
)
29+
30+
// Main is a main function for the external processor exposed
31+
// for allowing users to build their own external processor.
32+
func Main() {
33+
flag.Parse()
34+
35+
var level slog.Level
36+
if err := level.UnmarshalText([]byte(*logLevel)); err != nil {
37+
log.Fatalf("failed to unmarshal log level: %v", err)
38+
}
39+
l := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
40+
Level: level,
41+
}))
42+
43+
l.Info("starting external processor", slog.String("version", version.Version))
44+
45+
if *configPath == "" {
46+
log.Fatal("configPath must be provided")
47+
}
48+
49+
ctx, cancel := context.WithCancel(context.Background())
50+
signalsChan := make(chan os.Signal, 1)
51+
signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
52+
go func() {
53+
<-signalsChan
54+
cancel()
55+
}()
56+
57+
// TODO: unix domain socket support.
58+
lis, err := net.Listen("tcp", *extProcPort)
59+
if err != nil {
60+
log.Fatalf("failed to listen: %v", err)
61+
}
62+
63+
server, err := extproc.NewServer[*extproc.Processor](l, extproc.NewProcessor)
64+
if err != nil {
65+
log.Fatalf("failed to create external processor server: %v", err)
66+
}
67+
68+
if err := extproc.StartConfigWatcher(ctx, *configPath, server, l, time.Second*5); err != nil {
69+
log.Fatalf("failed to start config watcher: %v", err)
70+
}
71+
72+
s := grpc.NewServer()
73+
extprocv3.RegisterExternalProcessorServer(s, server)
74+
grpc_health_v1.RegisterHealthServer(s, server)
75+
go func() {
76+
<-ctx.Done()
77+
s.GracefulStop()
78+
}()
79+
_ = s.Serve(lis)
80+
}

‎examples/extproc_custom_router/README.md

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
This example shows how to insert a custom router in the custom external process using `filterconfig` package.

‎examples/extproc_custom_router/main.go

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
6+
"github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
7+
"github.com/envoyproxy/ai-gateway/extprocapi"
8+
"github.com/envoyproxy/ai-gateway/filterconfig"
9+
)
10+
11+
// newCustomRouter implements [extprocapi.NewCustomRouter].
12+
func newCustomRouter(defaultRouter extprocapi.Router, config *filterconfig.Config) extprocapi.Router {
13+
// You can poke the current configuration of the routes, and the list of backends
14+
// specified in the AIGatewayRoute.Rules, etc.
15+
return &myCustomRouter{config: config, defaultRouter: defaultRouter}
16+
}
17+
18+
// myCustomRouter implements [extprocapi.Router].
19+
type myCustomRouter struct {
20+
config *filterconfig.Config
21+
defaultRouter extprocapi.Router
22+
}
23+
24+
// Calculate implements [extprocapi.Router.Calculate].
25+
func (m *myCustomRouter) Calculate(headers map[string]string) (backend *filterconfig.Backend, err error) {
26+
// Simply logs the headers and delegates the calculation to the default router.
27+
modelName, ok := headers[m.config.ModelNameHeaderKey]
28+
if !ok {
29+
panic("model name not found in the headers")
30+
}
31+
fmt.Printf("model name: %s\n", modelName)
32+
return m.defaultRouter.Calculate(headers)
33+
}
34+
35+
// This demonstrates how to build a custom router for the external processor.
36+
func main() {
37+
// Initializes the custom router.
38+
extprocapi.NewCustomRouter = newCustomRouter
39+
// Executes the main function of the external processor.
40+
mainlib.Main()
41+
}

‎extprocapi/exptorcapi.go

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
// Package extprocapi is for building a custom external process.
2+
package extprocapi
3+
4+
import "github.com/envoyproxy/ai-gateway/filterconfig"
5+
6+
// NewCustomRouter is the function to create a custom router over the default router.
7+
// This is nil by default and can be set by the custom build of external processor.
8+
var NewCustomRouter NewCustomRouterFn
9+
10+
// NewCustomRouterFn is the function signature for [NewCustomRouter].
11+
//
12+
// It accepts the exptproc config passed to the AI Gateway filter and returns a [Router].
13+
// This is called when the new configuration is loaded.
14+
//
15+
// The defaultRouter can be used to delegate the calculation to the default router implementation.
16+
type NewCustomRouterFn func(defaultRouter Router, config *filterconfig.Config) Router
17+
18+
// Router is the interface for the router.
19+
//
20+
// Router must be goroutine-safe as it is shared across multiple requests.
21+
type Router interface {
22+
// Calculate determines the backend to route to based on the request headers.
23+
//
24+
// The request headers include the populated [filterconfig.Config.ModelNameHeaderKey]
25+
// with the parsed model name based on the [filterconfig.Config] given to the NewCustomRouterFn.
26+
//
27+
// Returns the backend.
28+
Calculate(requestHeaders map[string]string) (backend *filterconfig.Backend, err error)
29+
}

‎internal/extproc/mocks_test.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/stretchr/testify/require"
1212
"google.golang.org/grpc/metadata"
1313

14+
"github.com/envoyproxy/ai-gateway/extprocapi"
1415
"github.com/envoyproxy/ai-gateway/filterconfig"
1516
"github.com/envoyproxy/ai-gateway/internal/extproc/router"
1617
"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
@@ -19,7 +20,7 @@ import (
1920
var (
2021
_ ProcessorIface = &mockProcessor{}
2122
_ translator.Translator = &mockTranslator{}
22-
_ router.Router = &mockRouter{}
23+
_ extprocapi.Router = &mockRouter{}
2324
)
2425

2526
func newMockProcessor(_ *processorConfig) *mockProcessor {

‎internal/extproc/processor.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1313
"google.golang.org/protobuf/types/known/structpb"
1414

15+
"github.com/envoyproxy/ai-gateway/extprocapi"
1516
"github.com/envoyproxy/ai-gateway/filterconfig"
1617
"github.com/envoyproxy/ai-gateway/internal/extproc/backendauth"
1718
"github.com/envoyproxy/ai-gateway/internal/extproc/router"
@@ -22,7 +23,7 @@ import (
2223
// This will be created by the server and passed to the processor when it detects a new configuration.
2324
type processorConfig struct {
2425
bodyParser router.RequestBodyParser
25-
router router.Router
26+
router extprocapi.Router
2627
ModelNameHeaderKey, selectedBackendHeaderKey string
2728
factories map[filterconfig.VersionedAPISchema]translator.Factory
2829
backendAuthHandlers map[string]backendauth.Handler

‎internal/extproc/router/router.go

+11-12
Original file line numberDiff line numberDiff line change
@@ -6,28 +6,27 @@ import (
66

77
"golang.org/x/exp/rand"
88

9+
"github.com/envoyproxy/ai-gateway/extprocapi"
910
"github.com/envoyproxy/ai-gateway/filterconfig"
1011
)
1112

12-
// Router is the interface for the router.
13-
type Router interface {
14-
// Calculate determines the backend to route to based on the headers.
15-
// Returns the backend name and the output schema.
16-
Calculate(headers map[string]string) (backend *filterconfig.Backend, err error)
17-
}
18-
19-
// router implements [Router].
13+
// router implements [extprocapi.Router].
2014
type router struct {
2115
rules []filterconfig.RouteRule
2216
rng *rand.Rand
2317
}
2418

25-
// NewRouter creates a new [Router] implementation for the given config.
26-
func NewRouter(config *filterconfig.Config) (Router, error) {
27-
return &router{rules: config.Rules, rng: rand.New(rand.NewSource(uint64(time.Now().UnixNano())))}, nil
19+
// NewRouter creates a new [extprocapi.Router] implementation for the given config.
20+
func NewRouter(config *filterconfig.Config, newCustomFn extprocapi.NewCustomRouterFn) (extprocapi.Router, error) {
21+
r := &router{rules: config.Rules, rng: rand.New(rand.NewSource(uint64(time.Now().UnixNano())))}
22+
if newCustomFn != nil {
23+
customRouter := newCustomFn(r, config)
24+
return customRouter, nil
25+
}
26+
return r, nil
2827
}
2928

30-
// Calculate implements [Router.Calculate].
29+
// Calculate implements [extprocapi.Router.Calculate].
3130
func (r *router) Calculate(headers map[string]string) (backend *filterconfig.Backend, err error) {
3231
var rule *filterconfig.RouteRule
3332
for i := range r.rules {

‎internal/extproc/router/router_test.go

+27-2
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,34 @@ import (
55

66
"github.com/stretchr/testify/require"
77

8+
"github.com/envoyproxy/ai-gateway/extprocapi"
89
"github.com/envoyproxy/ai-gateway/filterconfig"
910
)
1011

12+
// dummyCustomRouter implements [extprocapi.Router].
13+
type dummyCustomRouter struct{ called bool }
14+
15+
func (c *dummyCustomRouter) Calculate(map[string]string) (*filterconfig.Backend, error) {
16+
c.called = true
17+
return nil, nil
18+
}
19+
20+
func TestRouter_NewRouter_Custom(t *testing.T) {
21+
r, err := NewRouter(&filterconfig.Config{}, func(defaultRouter extprocapi.Router, config *filterconfig.Config) extprocapi.Router {
22+
require.NotNil(t, defaultRouter)
23+
_, ok := defaultRouter.(*router)
24+
require.True(t, ok) // Checking if the default router is correctly passed.
25+
return &dummyCustomRouter{}
26+
})
27+
require.NoError(t, err)
28+
_, ok := r.(*dummyCustomRouter)
29+
require.True(t, ok)
30+
31+
_, err = r.Calculate(nil)
32+
require.NoError(t, err)
33+
require.True(t, r.(*dummyCustomRouter).called)
34+
}
35+
1136
func TestRouter_Calculate(t *testing.T) {
1237
outSchema := filterconfig.VersionedAPISchema{Schema: filterconfig.APISchemaOpenAI}
1338
_r, err := NewRouter(&filterconfig.Config{
@@ -30,7 +55,7 @@ func TestRouter_Calculate(t *testing.T) {
3055
},
3156
},
3257
},
33-
})
58+
}, nil)
3459
require.NoError(t, err)
3560
r, ok := _r.(*router)
3661
require.True(t, ok)
@@ -62,7 +87,7 @@ func TestRouter_Calculate(t *testing.T) {
6287
}
6388

6489
func TestRouter_selectBackendFromRule(t *testing.T) {
65-
_r, err := NewRouter(&filterconfig.Config{})
90+
_r, err := NewRouter(&filterconfig.Config{}, nil)
6691
require.NoError(t, err)
6792
r, ok := _r.(*router)
6893
require.True(t, ok)

‎internal/extproc/server.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"google.golang.org/grpc/health/grpc_health_v1"
1313
"google.golang.org/grpc/status"
1414

15+
"github.com/envoyproxy/ai-gateway/extprocapi"
1516
"github.com/envoyproxy/ai-gateway/filterconfig"
1617
"github.com/envoyproxy/ai-gateway/internal/extproc/backendauth"
1718
"github.com/envoyproxy/ai-gateway/internal/extproc/router"
@@ -37,7 +38,7 @@ func (s *Server[P]) LoadConfig(config *filterconfig.Config) error {
3738
if err != nil {
3839
return fmt.Errorf("cannot create request body parser: %w", err)
3940
}
40-
rt, err := router.NewRouter(config)
41+
rt, err := router.NewRouter(config, extprocapi.NewCustomRouter)
4142
if err != nil {
4243
return fmt.Errorf("cannot create router: %w", err)
4344
}

0 commit comments

Comments
 (0)