Repository: envoyproxy/ai-gateway
Pull request #425 (2025-02-25 21:26:39)

extproc: renames translator interface
**Commit Message**

After #325, the translator is not where the abstraction over :path is implemented, but it moved to the choice of processor. As a result, the translator interface has become effectively tied to a specific endpoint, notably /v1/chat/completion. This renames the `translator.Translator` interface accordingly not only to remove the unnecessary code path but also facilitates the implementation of additional features like metrics. 

**Related Issues/PRs (if applicable)**

Follow up on #325 and contribute to #316

mathetake: FYI @nacx 

Patch:
diff --git a/internal/extproc/chatcompletion_processor.go b/internal/extproc/chatcompletion_processor.go
index 9b4e5c3..2613173 100644
--- a/internal/extproc/chatcompletion_processor.go
+++ b/internal/extproc/chatcompletion_processor.go
@@ -46,7 +46,7 @@ type chatCompletionProcessor struct {
 	requestHeaders   map[string]string
 	responseHeaders  map[string]string
 	responseEncoding string
-	translator       translator.Translator
+	translator       translator.OpenAIChatCompletionTranslator
 	// cost is the cost of the request that is accumulated during the processing of the response.
 	costs translator.LLMTokenUsage
 }
@@ -213,7 +213,7 @@ func (c *chatCompletionProcessor) ProcessResponseBody(_ context.Context, body *e
 	return resp, nil
 }
 
-func parseOpenAIChatCompletionBody(body *extprocv3.HttpBody) (modelName string, rb translator.RequestBody, err error) {
+func parseOpenAIChatCompletionBody(body *extprocv3.HttpBody) (modelName string, rb *openai.ChatCompletionRequest, err error) {
 	var openAIReq openai.ChatCompletionRequest
 	if err := json.Unmarshal(body.Body, &openAIReq); err != nil {
 		return "", nil, fmt.Errorf("failed to unmarshal body: %w", err)
diff --git a/internal/extproc/translator/openai_awsbedrock.go b/internal/extproc/translator/openai_awsbedrock.go
index 01e1e8f..15df1a4 100644
--- a/internal/extproc/translator/openai_awsbedrock.go
+++ b/internal/extproc/translator/openai_awsbedrock.go
@@ -26,7 +26,7 @@ import (
 )
 
 // NewChatCompletionOpenAIToAWSBedrockTranslator implements [Factory] for OpenAI to AWS Bedrock translation.
-func NewChatCompletionOpenAIToAWSBedrockTranslator() Translator {
+func NewChatCompletionOpenAIToAWSBedrockTranslator() OpenAIChatCompletionTranslator {
 	return &openAIToAWSBedrockTranslatorV1ChatCompletion{}
 }
 
@@ -41,14 +41,9 @@ type openAIToAWSBedrockTranslatorV1ChatCompletion struct {
 }
 
 // RequestBody implements [Translator.RequestBody].
-func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) RequestBody(body RequestBody) (
+func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) RequestBody(openAIReq *openai.ChatCompletionRequest) (
 	headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, override *extprocv3http.ProcessingMode, err error,
 ) {
-	openAIReq, ok := body.(*openai.ChatCompletionRequest)
-	if !ok {
-		return nil, nil, nil, fmt.Errorf("unexpected body type: %T", body)
-	}
-
 	var pathTemplate string
 	if openAIReq.Stream {
 		o.stream = true
diff --git a/internal/extproc/translator/openai_openai.go b/internal/extproc/translator/openai_openai.go
index 62c110b..7052e12 100644
--- a/internal/extproc/translator/openai_openai.go
+++ b/internal/extproc/translator/openai_openai.go
@@ -19,7 +19,7 @@ import (
 )
 
 // NewChatCompletionOpenAIToOpenAITranslator implements [Factory] for OpenAI to OpenAI translation.
-func NewChatCompletionOpenAIToOpenAITranslator() Translator {
+func NewChatCompletionOpenAIToOpenAITranslator() OpenAIChatCompletionTranslator {
 	return &openAIToOpenAITranslatorV1ChatCompletion{}
 }
 
@@ -31,13 +31,9 @@ type openAIToOpenAITranslatorV1ChatCompletion struct {
 }
 
 // RequestBody implements [Translator.RequestBody].
-func (o *openAIToOpenAITranslatorV1ChatCompletion) RequestBody(body RequestBody) (
+func (o *openAIToOpenAITranslatorV1ChatCompletion) RequestBody(req *openai.ChatCompletionRequest) (
 	headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, override *extprocv3http.ProcessingMode, err error,
 ) {
-	req, ok := body.(*openai.ChatCompletionRequest)
-	if !ok {
-		return nil, nil, nil, fmt.Errorf("unexpected body type: %T", body)
-	}
 	if req.Stream {
 		o.stream = true
 		override = &extprocv3http.ProcessingMode{
diff --git a/internal/extproc/translator/translator.go b/internal/extproc/translator/translator.go
index 09f872e..9838520 100644
--- a/internal/extproc/translator/translator.go
+++ b/internal/extproc/translator/translator.go
@@ -12,6 +12,8 @@ import (
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3http "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/http/ext_proc/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+
+	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 )
 
 var (
@@ -30,21 +32,16 @@ func isGoodStatusCode(code int) bool {
 	return code >= 200 && code < 300
 }
 
-// RequestBody is the union of all request body types. TODO: maybe we should just define Translator interface per endpoint.
-type RequestBody any
-
-// Translator translates the request and response messages between the client and the backend API schemas for a specific path.
-// The implementation can embed [defaultTranslator] to avoid implementing all methods.
-//
-// The instance of [Translator] is created by a [Factory].
+// OpenAIChatCompletionTranslator translates the request and response messages between the client and the backend API schemas
+// for /v1/chat/completion endpoint of OpenAI.
 //
 // This is created per request and is not thread-safe.
-type Translator interface {
+type OpenAIChatCompletionTranslator interface {
 	// RequestBody translates the request body.
-	// 	- `body` is the request body already parsed by [router.RequestBodyParser]. The concrete type is specific to the schema and the path.
+	// 	- `body` is the request body parsed into the [openai.ChatCompletionRequest].
 	//	- This returns `headerMutation` and `bodyMutation` that can be nil to indicate no mutation.
 	//  - This returns `override` that to change the processing mode. This is used to process streaming requests properly.
-	RequestBody(body RequestBody) (
+	RequestBody(body *openai.ChatCompletionRequest) (
 		headerMutation *extprocv3.HeaderMutation,
 		bodyMutation *extprocv3.BodyMutation,
 		override *extprocv3http.ProcessingMode,


Test patch:
diff --git a/internal/extproc/mocks_test.go b/internal/extproc/mocks_test.go
index 482371c..6b94e28 100644
--- a/internal/extproc/mocks_test.go
+++ b/internal/extproc/mocks_test.go
@@ -19,13 +19,14 @@ import (
 
 	"github.com/envoyproxy/ai-gateway/filterapi"
 	"github.com/envoyproxy/ai-gateway/filterapi/x"
+	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
 )
 
 var (
-	_ Processor             = &mockProcessor{}
-	_ translator.Translator = &mockTranslator{}
-	_ x.Router              = &mockRouter{}
+	_ Processor                                 = &mockProcessor{}
+	_ translator.OpenAIChatCompletionTranslator = &mockTranslator{}
+	_ x.Router                                  = &mockRouter{}
 )
 
 func newMockProcessor(_ *processorConfig, _ *slog.Logger) Processor {
@@ -69,7 +70,7 @@ func (m mockProcessor) ProcessResponseBody(_ context.Context, body *extprocv3.Ht
 type mockTranslator struct {
 	t                 *testing.T
 	expHeaders        map[string]string
-	expRequestBody    translator.RequestBody
+	expRequestBody    *openai.ChatCompletionRequest
 	expResponseBody   *extprocv3.HttpBody
 	retHeaderMutation *extprocv3.HeaderMutation
 	retBodyMutation   *extprocv3.BodyMutation
@@ -78,19 +79,19 @@ type mockTranslator struct {
 	retErr            error
 }
 
-// RequestBody implements [translator.Translator.RequestBody].
-func (m mockTranslator) RequestBody(body translator.RequestBody) (headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, override *extprocv3http.ProcessingMode, err error) {
+// RequestBody implements [translator.OpenAIChatCompletionTranslator].
+func (m mockTranslator) RequestBody(body *openai.ChatCompletionRequest) (headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, override *extprocv3http.ProcessingMode, err error) {
 	require.Equal(m.t, m.expRequestBody, body)
 	return m.retHeaderMutation, m.retBodyMutation, m.retOverride, m.retErr
 }
 
-// ResponseHeaders implements [translator.Translator.ResponseHeaders].
+// ResponseHeaders implements [translator.OpenAIChatCompletionTranslator].
 func (m mockTranslator) ResponseHeaders(headers map[string]string) (headerMutation *extprocv3.HeaderMutation, err error) {
 	require.Equal(m.t, m.expHeaders, headers)
 	return m.retHeaderMutation, m.retErr
 }
 
-// ResponseError implements [translator.Translator.ResponseError].
+// ResponseError implements [translator.OpenAIChatCompletionTranslator].
 func (m mockTranslator) ResponseError(_ map[string]string, body io.Reader) (headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, err error) {
 	if m.expResponseBody != nil {
 		buf, err := io.ReadAll(body)
@@ -100,7 +101,7 @@ func (m mockTranslator) ResponseError(_ map[string]string, body io.Reader) (head
 	return m.retHeaderMutation, m.retBodyMutation, m.retErr
 }
 
-// ResponseBody implements [translator.Translator.ResponseBody].
+// ResponseBody implements [translator.OpenAIChatCompletionTranslator].
 func (m mockTranslator) ResponseBody(_ map[string]string, body io.Reader, _ bool) (headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, tokenUsage translator.LLMTokenUsage, err error) {
 	if m.expResponseBody != nil {
 		buf, err := io.ReadAll(body)
diff --git a/internal/extproc/translator/openai_awsbedrock_test.go b/internal/extproc/translator/openai_awsbedrock_test.go
index e57f27f..d4a7f6f 100644
--- a/internal/extproc/translator/openai_awsbedrock_test.go
+++ b/internal/extproc/translator/openai_awsbedrock_test.go
@@ -27,11 +27,6 @@ import (
 )
 
 func TestOpenAIToAWSBedrockTranslatorV1ChatCompletion_RequestBody(t *testing.T) {
-	t.Run("invalid body", func(t *testing.T) {
-		o := &openAIToAWSBedrockTranslatorV1ChatCompletion{}
-		_, _, _, err := o.RequestBody(&extprocv3.HttpBody{Body: []byte("invalid")})
-		require.Error(t, err)
-	})
 	tests := []struct {
 		name   string
 		output awsbedrock.ConverseInput
@@ -665,7 +660,7 @@ func TestOpenAIToAWSBedrockTranslatorV1ChatCompletion_RequestBody(t *testing.T)
 		t.Run(tt.name, func(t *testing.T) {
 			o := &openAIToAWSBedrockTranslatorV1ChatCompletion{}
 			originalReq := tt.input
-			hm, bm, mode, err := o.RequestBody(RequestBody(&originalReq))
+			hm, bm, mode, err := o.RequestBody(&originalReq)
 			var expPath string
 			if tt.input.Stream {
 				expPath = fmt.Sprintf("/model/%s/converse-stream", tt.input.Model)
diff --git a/internal/extproc/translator/openai_openai_test.go b/internal/extproc/translator/openai_openai_test.go
index 05d0595..403141d 100644
--- a/internal/extproc/translator/openai_openai_test.go
+++ b/internal/extproc/translator/openai_openai_test.go
@@ -23,18 +23,13 @@ import (
 )
 
 func TestOpenAIToOpenAITranslatorV1ChatCompletionRequestBody(t *testing.T) {
-	t.Run("invalid body", func(t *testing.T) {
-		o := &openAIToOpenAITranslatorV1ChatCompletion{}
-		_, _, _, err := o.RequestBody(&extprocv3.HttpBody{Body: []byte("invalid")})
-		require.Error(t, err)
-	})
 	t.Run("valid body", func(t *testing.T) {
 		for _, stream := range []bool{true, false} {
 			t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
 				originalReq := &openai.ChatCompletionRequest{Model: "foo-bar-ai", Stream: stream}
 
 				o := &openAIToOpenAITranslatorV1ChatCompletion{}
-				hm, bm, mode, err := o.RequestBody(RequestBody(originalReq))
+				hm, bm, mode, err := o.RequestBody(originalReq)
 				require.Nil(t, bm)
 				require.NoError(t, err)
 				require.Equal(t, stream, o.stream)