feat:update

Author: zhouyy

assistant.go

@@ -76,6 +76,9 @@ type AssistantRequest struct {
 	ResponseFormat any                    `json:"response_format,omitempty"`
 	Temperature    *float32               `json:"temperature,omitempty"`
 	TopP           *float32               `json:"top_p,omitempty"`
+	ExtraHeaders   map[string]string      `json:"extra_headers,omitempty"`
+	ExtraQuery     map[string]string      `json:"extra_query,omitempty"`
+	ExtraBody      map[string]any         `json:"extra_body,omitempty"`
 }
 
 // MarshalJSON provides a custom marshaller for the assistant request to handle the API use cases
@@ -137,7 +140,11 @@ type AssistantFilesList struct {
 // CreateAssistant creates a new assistant.
 func (c *Client) CreateAssistant(ctx context.Context, request AssistantRequest) (response Assistant, err error) {
 	req, err := c.newRequest(ctx, http.MethodPost, c.fullURL(assistantsSuffix), withBody(request),
-		withBetaAssistantVersion(c.config.AssistantVersion))
+		withBetaAssistantVersion(c.config.AssistantVersion),
+		withExtraHeaders(request.ExtraHeaders),
+		withExtraQuery(request.ExtraQuery),
+		withExtraBody(request.ExtraBody),
+	)
 	if err != nil {
 		return
 	}

audio.go

@@ -49,6 +49,9 @@ type AudioRequest struct {
 	Language               string // Only for transcription.
 	Format                 AudioResponseFormat
 	TimestampGranularities []TranscriptionTimestampGranularity // Only for transcription.
+	ExtraHeaders           map[string]string
+	ExtraQuery             map[string]string
+	ExtraBody              map[string]any
 }
 
 // AudioResponse represents a response structure for audio API.
@@ -128,6 +131,9 @@ func (c *Client) callAudioAPI(
 		c.fullURL(urlSuffix, withModel(request.Model)),
 		withBody(&formBody),
 		withContentType(builder.FormDataContentType()),
+		withExtraHeaders(request.ExtraHeaders),
+		withExtraQuery(request.ExtraQuery),
+		withExtraBody(request.ExtraBody),
 	)
 	if err != nil {
 		return AudioResponse{}, err

batch.go

@@ -97,10 +97,13 @@ type BatchRequestCounts struct {
 }
 
 type CreateBatchRequest struct {
-	InputFileID      string         `json:"input_file_id"`
-	Endpoint         BatchEndpoint  `json:"endpoint"`
-	CompletionWindow string         `json:"completion_window"`
-	Metadata         map[string]any `json:"metadata"`
+	InputFileID      string            `json:"input_file_id"`
+	Endpoint         BatchEndpoint     `json:"endpoint"`
+	CompletionWindow string            `json:"completion_window"`
+	Metadata         map[string]any    `json:"metadata"`
+	ExtraHeaders     map[string]string `json:"extra_headers"`
+	ExtraQuery       map[string]string `json:"extra_query"`
+	ExtraBody        map[string]any    `json:"extra_body"`
 }
 
 type BatchResponse struct {
@@ -117,7 +120,12 @@ func (c *Client) CreateBatch(
 		request.CompletionWindow = "24h"
 	}
 
-	req, err := c.newRequest(ctx, http.MethodPost, c.fullURL(batchesSuffix), withBody(request))
+	req, err := c.newRequest(ctx, http.MethodPost, c.fullURL(batchesSuffix),
+		withBody(request),
+		withExtraHeaders(request.ExtraHeaders),
+		withExtraQuery(request.ExtraQuery),
+		withExtraBody(request.ExtraBody),
+	)
 	if err != nil {
 		return
 	}

chat.go

@@ -265,43 +265,10 @@ type ChatCompletionRequest struct {
 	Metadata map[string]string `json:"metadata,omitempty"`
 	// Add additional JSON properties to the request
 	ExtraBody map[string]any `json:"extra_body,omitempty"`
-}
-
-func (m ChatCompletionRequest) MarshalJSON() ([]byte, error) {
-	// Create a new anonymous struct that omits ExtraBody
-	type Alias ChatCompletionRequest
-	temp := struct {
-		Alias
-		ExtraBody map[string]any `json:"-"` // Omit ExtraBody from direct serialization
-	}{
-		Alias:     Alias(m),
-		ExtraBody: m.ExtraBody,
-	}
-
-	// First marshal the main structure
-	data, err := json.Marshal(temp)
-	if err != nil {
-		return nil, err
-	}
-
-	// If there's no ExtraBody, return the marshaled data as is
-	if len(m.ExtraBody) == 0 {
-		return data, nil
-	}
-
-	// Unmarshal into a map to modify the JSON structure
-	var rawMap map[string]any
-	if err := json.Unmarshal(data, &rawMap); err != nil {
-		return nil, err
-	}
-
-	// Add ExtraBody fields to the root level
-	for k, v := range m.ExtraBody {
-		rawMap[k] = v
-	}
-
-	// Marshal the combined map back to JSON
-	return json.Marshal(rawMap)
+	// ExtraHeaders to add to the request
+	ExtraHeaders map[string]string `json:"extra_headers,omitempty"`
+	// ExtraQueryParams to add to the request
+	ExtraQuery map[string]string `json:"extra_query,omitempty"`
 }
 
 type StreamOptions struct {
@@ -442,6 +409,9 @@ func (c *Client) CreateChatCompletion(
 		http.MethodPost,
 		c.fullURL(urlSuffix, withModel(request.Model)),
 		withBody(request),
+		withExtraHeaders(request.ExtraHeaders),
+		withExtraQuery(request.ExtraQuery),
+		withExtraBody(request.ExtraBody),
 	)
 	if err != nil {
 		return

client.go

@@ -72,8 +72,11 @@ func NewOrgClient(authToken, org string) *Client {
 }
 
 type requestOptions struct {
-	body   any
-	header http.Header
+	body         any
+	header       http.Header
+	extraHeaders map[string]string
+	extraQuery   map[string]string
+	extraBody    map[string]any
 }
 
 type requestOption func(*requestOptions)
@@ -84,6 +87,30 @@ func withBody(body any) requestOption {
 	}
 }
 
+func withHeader(header http.Header) requestOption {
+	return func(args *requestOptions) {
+		args.header = header
+	}
+}
+
+func withExtraHeaders(extraHeaders map[string]string) requestOption {
+	return func(args *requestOptions) {
+		args.extraHeaders = extraHeaders
+	}
+}
+
+func withExtraQuery(extraQuery map[string]string) requestOption {
+	return func(args *requestOptions) {
+		args.extraQuery = extraQuery
+	}
+}
+
+func withExtraBody(extraBody map[string]any) requestOption {
+	return func(args *requestOptions) {
+		args.extraBody = extraBody
+	}
+}
+
 func withContentType(contentType string) requestOption {
 	return func(args *requestOptions) {
 		args.header.Set("Content-Type", contentType)
@@ -105,7 +132,15 @@ func (c *Client) newRequest(ctx context.Context, method, url string, setters ...
 	for _, setter := range setters {
 		setter(args)
 	}
-	req, err := c.requestBuilder.Build(ctx, method, url, args.body, args.header)
+	req, err := c.requestBuilder.Build(ctx, &utils.Request{
+		Method:       method,
+		URL:          url,
+		Body:         args.body,
+		Header:       args.header,
+		ExtraHeaders: args.extraHeaders,
+		ExtraQuery:   args.extraQuery,
+		ExtraBody:    args.extraBody,
+	})
 	if err != nil {
 		return nil, err
 	}

client_test.go

@@ -10,6 +10,7 @@ import (
 	"reflect"
 	"testing"
 
+	utils "github.com/sashabaranov/go-openai/internal"
 	"github.com/sashabaranov/go-openai/internal/test"
 	"github.com/sashabaranov/go-openai/internal/test/checks"
 )
@@ -18,7 +19,7 @@ var errTestRequestBuilderFailed = errors.New("test request builder failed")
 
 type failingRequestBuilder struct{}
 
-func (*failingRequestBuilder) Build(_ context.Context, _, _ string, _ any, _ http.Header) (*http.Request, error) {
+func (*failingRequestBuilder) Build(_ context.Context, _ *utils.Request) (*http.Request, error) {
 	return nil, errTestRequestBuilderFailed
 }
 

embeddings.go

@@ -159,7 +159,10 @@ type EmbeddingRequest struct {
 	EncodingFormat EmbeddingEncodingFormat `json:"encoding_format,omitempty"`
 	// Dimensions The number of dimensions the resulting output embeddings should have.
 	// Only supported in text-embedding-3 and later models.
-	Dimensions int `json:"dimensions,omitempty"`
+	Dimensions   int               `json:"dimensions,omitempty"`
+	ExtraHeaders map[string]string `json:"extra_headers,omitempty"`
+	ExtraQuery   map[string]string `json:"extra_query,omitempty"`
+	ExtraBody    map[string]any    `json:"extra_body,omitempty"`
 }
 
 func (r EmbeddingRequest) Convert() EmbeddingRequest {
@@ -246,6 +249,9 @@ func (c *Client) CreateEmbeddings(
 		http.MethodPost,
 		c.fullURL("/embeddings", withModel(string(baseReq.Model))),
 		withBody(baseReq),
+		withExtraHeaders(baseReq.ExtraHeaders),
+		withExtraQuery(baseReq.ExtraQuery),
+		withExtraBody(baseReq.ExtraBody),
 	)
 	if err != nil {
 		return

internal/marshaller.go

@@ -6,10 +6,15 @@ import (
 
 type Marshaller interface {
 	Marshal(value any) ([]byte, error)
+	Unmarshal(data []byte, value any) error
 }
 
 type JSONMarshaller struct{}
 
 func (jm *JSONMarshaller) Marshal(value any) ([]byte, error) {
 	return json.Marshal(value)
 }
+
+func (jm *JSONMarshaller) Unmarshal(data []byte, value any) error {
+	return json.Unmarshal(data, value)
+}

internal/request_builder.go

@@ -3,12 +3,25 @@ package openai
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"net/http"
+	"net/url"
+	"strings"
 )
 
+type Request struct {
+	Method       string
+	URL          string
+	Body         any
+	Header       http.Header
+	ExtraBody    map[string]any
+	ExtraHeaders map[string]string
+	ExtraQuery   map[string]string
+}
+
 type RequestBuilder interface {
-	Build(ctx context.Context, method, url string, body any, header http.Header) (*http.Request, error)
+	Build(ctx context.Context, request *Request) (*http.Request, error)
 }
 
 type HTTPRequestBuilder struct {
@@ -23,30 +36,61 @@ func NewRequestBuilder() *HTTPRequestBuilder {
 
 func (b *HTTPRequestBuilder) Build(
 	ctx context.Context,
-	method string,
-	url string,
-	body any,
-	header http.Header,
+	request *Request,
 ) (req *http.Request, err error) {
 	var bodyReader io.Reader
-	if body != nil {
-		if v, ok := body.(io.Reader); ok {
+	if request.Body != nil {
+		if v, ok := request.Body.(io.Reader); ok {
 			bodyReader = v
 		} else {
 			var reqBytes []byte
-			reqBytes, err = b.marshaller.Marshal(body)
+			reqBytes, err = b.marshaller.Marshal(request.Body)
 			if err != nil {
 				return
 			}
+
+			if request.ExtraBody != nil {
+				rawMap := make(map[string]any)
+				err = b.marshaller.Unmarshal(reqBytes, &rawMap)
+				if err != nil {
+					return
+				}
+
+				for k, v := range request.ExtraBody {
+					rawMap[k] = v
+				}
+				reqBytes, err = b.marshaller.Marshal(rawMap)
+				if err != nil {
+					return
+				}
+			}
+
 			bodyReader = bytes.NewBuffer(reqBytes)
 		}
 	}
-	req, err = http.NewRequestWithContext(ctx, method, url, bodyReader)
+
+	requestUrl := request.URL
+	if request.ExtraQuery != nil {
+		for k, v := range request.ExtraQuery {
+			requestUrl = fmt.Sprintf("%s&%s=%s", requestUrl, k, url.QueryEscape(v))
+		}
+	}
+
+	requestUrl = strings.TrimSuffix(requestUrl, "&")
+
+	req, err = http.NewRequestWithContext(ctx, request.Method, requestUrl, bodyReader)
 	if err != nil {
 		return
 	}
-	if header != nil {
-		req.Header = header
+	if request.Header != nil {
+		req.Header = request.Header
+	} else {
+		req.Header = make(http.Header)
 	}
+
+	for k, v := range request.ExtraHeaders {
+		req.Header.Set(k, v)
+	}
+
 	return
 }