Merge pull request #7 from meguminnnnnnnnn/feat/option

feat: add WithExtraHeader option

Author: mrh997

chat.go

@@ -510,9 +510,9 @@ func (c *Client) CreateChatCompletion(
 	}
 
 	body := any(request)
-	if ccOpts.RequestBodySetter != nil {
+	if ccOpts.RequestBodyModifier != nil {
 		var newBody io.Reader
-		newBody, err = c.getNewRequestBody(request, ccOpts.RequestBodySetter)
+		newBody, err = c.getNewRequestBody(request, ccOpts.RequestBodyModifier)
 		if err != nil {
 			return response, err
 		}
@@ -524,6 +524,7 @@ func (c *Client) CreateChatCompletion(
 		http.MethodPost,
 		c.fullURL(urlSuffix, withModel(request.Model)),
 		withBody(body),
+		withExtraHeader(ccOpts.ExtraHeader),
 	)
 	if err != nil {
 		return
@@ -533,15 +534,15 @@ func (c *Client) CreateChatCompletion(
 	return
 }
 
-func (c *Client) getNewRequestBody(request ChatCompletionRequest, setter RequestBodySetter) (io.Reader, error) {
+func (c *Client) getNewRequestBody(request ChatCompletionRequest, modifier RequestBodyModifier) (io.Reader, error) {
 	marshaller := openai.JSONMarshaller{}
 
 	body, err := marshaller.Marshal(request)
 	if err != nil {
 		return nil, err
 	}
 
-	newBody, err := setter(body)
+	newBody, err := modifier(body)
 	if err != nil {
 		return nil, err
 	}

chat_stream.go

@@ -99,9 +99,9 @@ func (c *Client) CreateChatCompletionStream(
 	}
 
 	body := any(request)
-	if ccOpts.RequestBodySetter != nil {
+	if ccOpts.RequestBodyModifier != nil {
 		var newBody io.Reader
-		newBody, err = c.getNewRequestBody(request, ccOpts.RequestBodySetter)
+		newBody, err = c.getNewRequestBody(request, ccOpts.RequestBodyModifier)
 		if err != nil {
 			return stream, err
 		}
@@ -113,6 +113,7 @@ func (c *Client) CreateChatCompletionStream(
 		http.MethodPost,
 		c.fullURL(urlSuffix, withModel(request.Model)),
 		withBody(body),
+		withExtraHeader(ccOpts.ExtraHeader),
 	)
 	if err != nil {
 		return nil, err

chat_test.go

@@ -52,12 +52,12 @@ func TestChatCompletionsWrongModel(t *testing.T) {
 	checks.ErrorIs(t, err, openai.ErrChatCompletionInvalidModel, msg)
 }
 
-func TestChatCompletionRequestWithRequestBodySetter(t *testing.T) {
+func TestChatCompletionRequestWithRequestBodyModifier(t *testing.T) {
 	client, server, teardown := setupOpenAITestServer()
 	defer teardown()
 	server.RegisterHandler("/v1/chat/completions", handleChatCompletionEndpoint)
 
-	opt := openai.WithRequestBodySetter(func(rawBody []byte) ([]byte, error) {
+	opt := openai.WithRequestBodyModifier(func(rawBody []byte) ([]byte, error) {
 		return rawBody, nil
 	})
 

client.go

@@ -98,6 +98,14 @@ func withExtraBody(extraBody map[string]any) requestOption {
 	}
 }
 
+func withExtraHeader(header map[string]string) requestOption {
+	return func(args *requestOptions) {
+		for k, v := range header {
+			args.header.Set(k, v)
+		}
+	}
+}
+
 func withContentType(contentType string) requestOption {
 	return func(args *requestOptions) {
 		args.header.Set("Content-Type", contentType)

option.go

@@ -1,15 +1,22 @@
 package openai
 
 type chatCompletionRequestOptions struct {
-	RequestBodySetter RequestBodySetter
+	RequestBodyModifier RequestBodyModifier
+	ExtraHeader         map[string]string
 }
 
 type ChatCompletionRequestOption func(*chatCompletionRequestOptions)
 
-type RequestBodySetter func(rawBody []byte) ([]byte, error)
+type RequestBodyModifier func(rawBody []byte) ([]byte, error)
 
-func WithRequestBodySetter(setter RequestBodySetter) ChatCompletionRequestOption {
+func WithRequestBodyModifier(modifier RequestBodyModifier) ChatCompletionRequestOption {
 	return func(opts *chatCompletionRequestOptions) {
-		opts.RequestBodySetter = setter
+		opts.RequestBodyModifier = modifier
+	}
+}
+
+func WithExtraHeader(header map[string]string) ChatCompletionRequestOption {
+	return func(opts *chatCompletionRequestOptions) {
+		opts.ExtraHeader = header
 	}
 }