Merge branch 'main' into main

Author: mathetake

internal/apischema/openai/openai.go

@@ -1167,9 +1167,9 @@ type ChatCompletionResponseChoiceMessage struct {
 	// Audio is the audio response generated by the model, if applicable.
 	Audio *ChatCompletionResponseChoiceMessageAudio `json:"audio,omitempty"`
 
-	// AWSBedRockResponseVendorFields is used to hold any non-standard fields from the backend,
+	// AWSBedrockResponseVendorFields is used to hold any non-standard fields from the backend,
 	// like "reasoningContent" from AWS Bedrock.
-	*AWSBedRockResponseVendorFields `json:",inline,omitempty"`
+	*AWSBedrockResponseVendorFields `json:",inline,omitempty"`
 }
 
 // URLCitation contains citation information for web search results.
@@ -1303,7 +1303,7 @@ type ChatCompletionResponseChunkChoiceDelta struct {
 	Role             string                               `json:"role,omitempty"`
 	ToolCalls        []ChatCompletionMessageToolCallParam `json:"tool_calls,omitempty"`
 	Annotations      *[]Annotation                        `json:"annotations,omitempty"`
-	ReasoningContent *AWSBedRockStreamReasoningContent    `json:"reasoning_content,omitempty"`
+	ReasoningContent *AWSBedrockStreamReasoningContent    `json:"reasoning_content,omitempty"`
 }
 
 // Error is described in the OpenAI API documentation
@@ -1516,20 +1516,20 @@ type AnthropicVendorFields struct {
 	Thinking *anthropic.ThinkingConfigParamUnion `json:"thinking,omitzero"`
 }
 
-// AWSBedRockResponseVendorFields contains extra vendor-specific fields to add to the openai response.
-type AWSBedRockResponseVendorFields struct {
+// AWSBedrockResponseVendorFields contains extra vendor-specific fields to add to the openai response.
+type AWSBedrockResponseVendorFields struct {
 	// Contains content regarding the reasoning that is carried out by the model. Reasoning refers to a Chain of Thought (CoT) that the model generates to enhance the accuracy of its final response.
 	// Note: This object is a Union. Only one member of this object can be specified or returned.
 	// Required: No
 	// See https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ReasoningContentBlock.html for more information.
-	ReasoningContent *AWSBedRockReasoningContent `json:"reasoning_content,omitzero"`
+	ReasoningContent *AWSBedrockReasoningContent `json:"reasoning_content,omitzero"`
 }
 
-type AWSBedRockReasoningContent struct {
+type AWSBedrockReasoningContent struct {
 	ReasoningContent *awsbedrock.ReasoningContentBlock `json:"reasoningContent,omitzero"`
 }
 
-type AWSBedRockStreamReasoningContent struct {
+type AWSBedrockStreamReasoningContent struct {
 	Text            string `json:"text,omitzero"`
 	Signature       string `json:"signature,omitzero"`
 	RedactedContent []byte `json:"redactedContent,omitzero"`

internal/extproc/chatcompletion_processor.go

@@ -6,12 +6,9 @@
 package extproc
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log/slog"
 	"strconv"
 
@@ -326,24 +323,18 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 			c.metrics.RecordRequestCompletion(ctx, true, c.requestHeaders)
 		}
 	}()
-	var br io.Reader
-	var isGzip bool
-	switch c.responseEncoding {
-	case "gzip":
-		br, err = gzip.NewReader(bytes.NewReader(body.Body))
-		if err != nil {
-			return nil, fmt.Errorf("failed to decode gzip: %w", err)
-		}
-		isGzip = true
-	default:
-		br = bytes.NewReader(body.Body)
+
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, c.responseEncoding)
+	if err != nil {
+		return nil, err
 	}
 
 	// Assume all responses have a valid status code header.
 	if code, _ := strconv.Atoi(c.responseHeaders[":status"]); !isGoodStatusCode(code) {
 		var headerMutation *extprocv3.HeaderMutation
 		var bodyMutation *extprocv3.BodyMutation
-		headerMutation, bodyMutation, err = c.translator.ResponseError(c.responseHeaders, br)
+		headerMutation, bodyMutation, err = c.translator.ResponseError(c.responseHeaders, decodingResult.reader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to transform response error: %w", err)
 		}
@@ -368,22 +359,13 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 		}, nil
 	}
 
-	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream, c.span)
+	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, decodingResult.reader, body.EndOfStream, c.span)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
-	if bodyMutation != nil && isGzip {
-		if headerMutation == nil {
-			headerMutation = &extprocv3.HeaderMutation{}
-		}
-		// TODO: this is a hotfix, we should update this to recompress since its in the header
-		// If the response was gzipped, ensure we remove the content-encoding header.
-		//
-		// This is only needed when the transformation is actually modifying the body. When the backend
-		// is in OpenAI format (and it's the first try before any retry), the response body is not modified,
-		// so we don't need to remove the header in that case.
-		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
-	}
+
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
 
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{

internal/extproc/embeddings_processor.go

@@ -6,12 +6,9 @@
 package extproc
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log/slog"
 	"strconv"
 
@@ -250,24 +247,18 @@ func (e *embeddingsProcessorUpstreamFilter) ProcessResponseBody(ctx context.Cont
 			e.metrics.RecordRequestCompletion(ctx, true, e.requestHeaders)
 		}
 	}()
-	var br io.Reader
-	var isGzip bool
-	switch e.responseEncoding {
-	case "gzip":
-		br, err = gzip.NewReader(bytes.NewReader(body.Body))
-		if err != nil {
-			return nil, fmt.Errorf("failed to decode gzip: %w", err)
-		}
-		isGzip = true
-	default:
-		br = bytes.NewReader(body.Body)
+
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, e.responseEncoding)
+	if err != nil {
+		return nil, err
 	}
 
 	// Assume all responses have a valid status code header.
 	if code, _ := strconv.Atoi(e.responseHeaders[":status"]); !isGoodStatusCode(code) {
 		var headerMutation *extprocv3.HeaderMutation
 		var bodyMutation *extprocv3.BodyMutation
-		headerMutation, bodyMutation, err = e.translator.ResponseError(e.responseHeaders, br)
+		headerMutation, bodyMutation, err = e.translator.ResponseError(e.responseHeaders, decodingResult.reader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to transform response error: %w", err)
 		}
@@ -285,22 +276,13 @@ func (e *embeddingsProcessorUpstreamFilter) ProcessResponseBody(ctx context.Cont
 		}, nil
 	}
 
-	headerMutation, bodyMutation, tokenUsage, err := e.translator.ResponseBody(e.responseHeaders, br, body.EndOfStream)
+	headerMutation, bodyMutation, tokenUsage, err := e.translator.ResponseBody(e.responseHeaders, decodingResult.reader, body.EndOfStream)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
-	if bodyMutation != nil && isGzip {
-		if headerMutation == nil {
-			headerMutation = &extprocv3.HeaderMutation{}
-		}
-		// TODO: this is a hotfix, we should update this to recompress since its in the header
-		// If the response was gzipped, ensure we remove the content-encoding header.
-		//
-		// This is only needed when the transformation is actually modifying the body. When the backend
-		// is in OpenAI format (and it's the first try before any retry), the response body is not modified,
-		// so we don't need to remove the header in that case.
-		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
-	}
+
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
 
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{

internal/extproc/messages_processor.go

@@ -6,7 +6,6 @@
 package extproc
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -134,6 +133,7 @@ type messagesProcessorUpstreamFilter struct {
 	config                 *processorConfig
 	requestHeaders         map[string]string
 	responseHeaders        map[string]string
+	responseEncoding       string
 	modelNameOverride      string
 	backendName            string
 	handler                backendauth.Handler
@@ -222,6 +222,9 @@ func (c *messagesProcessorUpstreamFilter) ProcessResponseHeaders(ctx context.Con
 	}()
 
 	c.responseHeaders = headersToMap(headers)
+	if enc := c.responseHeaders["content-encoding"]; enc != "" {
+		c.responseEncoding = enc
+	}
 	headerMutation, err := c.translator.ResponseHeaders(c.responseHeaders)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response headers: %w", err)
@@ -250,14 +253,21 @@ func (c *messagesProcessorUpstreamFilter) ProcessResponseBody(ctx context.Contex
 		}
 	}()
 
-	// Simple passthrough: just pass the body as-is without any complex handling.
-	br := bytes.NewReader(body.Body)
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, c.responseEncoding)
+	if err != nil {
+		return nil, err
+	}
 
-	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream)
+	// headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream).
+	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, decodingResult.reader, body.EndOfStream)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
 
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
+
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{
 			ResponseBody: &extprocv3.BodyResponse{

internal/extproc/translator/openai_awsbedrock.go

@@ -668,10 +668,10 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) ResponseBody(_ map[string
 				choice.Message.Content = output.Text
 			}
 		case output.ReasoningContent != nil:
-			if choice.Message.AWSBedRockResponseVendorFields == nil {
-				choice.Message.AWSBedRockResponseVendorFields = &openai.AWSBedRockResponseVendorFields{}
+			if choice.Message.AWSBedrockResponseVendorFields == nil {
+				choice.Message.AWSBedrockResponseVendorFields = &openai.AWSBedrockResponseVendorFields{}
 			}
-			choice.Message.ReasoningContent = &openai.AWSBedRockReasoningContent{ReasoningContent: output.ReasoningContent}
+			choice.Message.ReasoningContent = &openai.AWSBedrockReasoningContent{ReasoningContent: output.ReasoningContent}
 		}
 	}
 	openAIResp.Choices = append(openAIResp.Choices, choice)
@@ -761,7 +761,7 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) convertEvent(event *awsbe
 				},
 			})
 		case event.Delta.ReasoningContent != nil:
-			reasoningDelta := &openai.AWSBedRockStreamReasoningContent{}
+			reasoningDelta := &openai.AWSBedrockStreamReasoningContent{}
 
 			// Map all relevant fields from the Bedrock delta to our flattened OpenAI delta struct.
 			if event.Delta.ReasoningContent.ReasoningText != nil {

internal/extproc/translator/openai_awsbedrock_test.go

@@ -1603,8 +1603,8 @@ func TestOpenAIToAWSBedrockTranslatorV1ChatCompletion_ResponseBody(t *testing.T)
 						Message: openai.ChatCompletionResponseChoiceMessage{
 							Role:    awsbedrock.ConversationRoleAssistant,
 							Content: ptr.To("This is the final answer."),
-							AWSBedRockResponseVendorFields: &openai.AWSBedRockResponseVendorFields{
-								ReasoningContent: &openai.AWSBedRockReasoningContent{
+							AWSBedrockResponseVendorFields: &openai.AWSBedrockResponseVendorFields{
+								ReasoningContent: &openai.AWSBedrockReasoningContent{
 									ReasoningContent: &awsbedrock.ReasoningContentBlock{
 										ReasoningText: &awsbedrock.ReasoningTextBlock{
 											Text: "This is the model's thought process.",
@@ -1807,7 +1807,7 @@ func TestOpenAIToAWSBedrockTranslator_convertEvent(t *testing.T) {
 				Choices: []openai.ChatCompletionResponseChunkChoice{
 					{
 						Delta: &openai.ChatCompletionResponseChunkChoiceDelta{
-							ReasoningContent: &openai.AWSBedRockStreamReasoningContent{
+							ReasoningContent: &openai.AWSBedrockStreamReasoningContent{
 								Text: "thinking...",
 							},
 						},

internal/extproc/util.go

@@ -5,6 +5,60 @@
 
 package extproc
 
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+
+	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+)
+
+// contentDecodingResult contains the result of content decoding operation.
+type contentDecodingResult struct {
+	reader    io.Reader
+	isEncoded bool
+}
+
+// decodeContentIfNeeded decompresses the response body based on the content-encoding header.
+// Currently supports gzip encoding, but can be extended to support other encodings in the future.
+// Returns a reader for the (potentially decompressed) body and metadata about the encoding.
+func decodeContentIfNeeded(body []byte, contentEncoding string) (contentDecodingResult, error) {
+	switch contentEncoding {
+	case "gzip":
+		reader, err := gzip.NewReader(bytes.NewReader(body))
+		if err != nil {
+			return contentDecodingResult{}, fmt.Errorf("failed to decode gzip: %w", err)
+		}
+		return contentDecodingResult{
+			reader:    reader,
+			isEncoded: true,
+		}, nil
+	default:
+		return contentDecodingResult{
+			reader:    bytes.NewReader(body),
+			isEncoded: false,
+		}, nil
+	}
+}
+
+// removeContentEncodingIfNeeded removes the content-encoding header if the body was modified and was encoded.
+// This is needed when the transformation modifies the body content but the response was originally compressed.
+func removeContentEncodingIfNeeded(headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, isEncoded bool) *extprocv3.HeaderMutation {
+	if bodyMutation != nil && isEncoded {
+		if headerMutation == nil {
+			headerMutation = &extprocv3.HeaderMutation{}
+		}
+		// TODO: this is a hotfix, we should update this to recompress since its in the header
+		// If the upstream response was compressed and we decompressed it,
+		// ensure we remove the content-encoding header.
+		//
+		// This is only needed when the transformation is actually modifying the body.
+		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
+	}
+	return headerMutation
+}
+
 // isGoodStatusCode checks if the HTTP status code of the upstream response is successful.
 // The 2xx - Successful: The request is received by upstream and processed successfully.
 // https://developer.mozilla.org/en-US/docs/Web/HTTP/Status#successful_responses