Merge pull request #5 from meguminnnnnnnnn/feat/option

feat: add WithRequestBodySetter Option

Author: mrh997

chat.go

@@ -1,11 +1,15 @@
 package openai
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
+	"io"
 	"net/http"
 
+	openai "github.com/meguminnnnnnnnn/go-openai/internal"
+
 	"github.com/meguminnnnnnnnn/go-openai/jsonschema"
 )
 
@@ -482,6 +486,7 @@ type ChatCompletionResponse struct {
 func (c *Client) CreateChatCompletion(
 	ctx context.Context,
 	request ChatCompletionRequest,
+	opts ...ChatCompletionRequestOption,
 ) (response ChatCompletionResponse, err error) {
 	if request.Stream {
 		err = ErrChatCompletionStreamNotSupported
@@ -499,11 +504,26 @@ func (c *Client) CreateChatCompletion(
 		return
 	}
 
+	ccOpts := &chatCompletionRequestOptions{}
+	for _, opt := range opts {
+		opt(ccOpts)
+	}
+
+	body := any(request)
+	if ccOpts.RequestBodySetter != nil {
+		var newBody io.Reader
+		newBody, err = c.getNewRequestBody(request, ccOpts.RequestBodySetter)
+		if err != nil {
+			return response, err
+		}
+		body = newBody
+	}
+
 	req, err := c.newRequest(
 		ctx,
 		http.MethodPost,
 		c.fullURL(urlSuffix, withModel(request.Model)),
-		withBody(request),
+		withBody(body),
 	)
 	if err != nil {
 		return
@@ -512,3 +532,19 @@ func (c *Client) CreateChatCompletion(
 	err = c.sendRequest(req, &response)
 	return
 }
+
+func (c *Client) getNewRequestBody(request ChatCompletionRequest, setter RequestBodySetter) (io.Reader, error) {
+	marshaller := openai.JSONMarshaller{}
+
+	body, err := marshaller.Marshal(request)
+	if err != nil {
+		return nil, err
+	}
+
+	newBody, err := setter(body)
+	if err != nil {
+		return nil, err
+	}
+
+	return bytes.NewBuffer(newBody), nil
+}

chat_test.go

@@ -52,6 +52,28 @@ func TestChatCompletionsWrongModel(t *testing.T) {
 	checks.ErrorIs(t, err, openai.ErrChatCompletionInvalidModel, msg)
 }
 
+func TestChatCompletionRequestWithRequestBodySetter(t *testing.T) {
+	client, server, teardown := setupOpenAITestServer()
+	defer teardown()
+	server.RegisterHandler("/v1/chat/completions", handleChatCompletionEndpoint)
+
+	opt := openai.WithRequestBodySetter(func(rawBody []byte) ([]byte, error) {
+		return rawBody, nil
+	})
+
+	_, err := client.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{
+		Model:               openai.O1Preview,
+		MaxCompletionTokens: 1000,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    openai.ChatMessageRoleUser,
+				Content: "Hello!",
+			},
+		},
+	}, opt)
+	checks.NoError(t, err)
+}
+
 func TestO1ModelsChatCompletionsDeprecatedFields(t *testing.T) {
 	tests := []struct {
 		name          string

option.go

@@ -0,0 +1,15 @@
+package openai
+
+type chatCompletionRequestOptions struct {
+	RequestBodySetter RequestBodySetter
+}
+
+type ChatCompletionRequestOption func(*chatCompletionRequestOptions)
+
+type RequestBodySetter func(rawBody []byte) ([]byte, error)
+
+func WithRequestBodySetter(setter RequestBodySetter) ChatCompletionRequestOption {
+	return func(opts *chatCompletionRequestOptions) {
+		opts.RequestBodySetter = setter
+	}
+}