Merge pull request #244 from charmbracelet/add-extra-body

Kujtim Hoxha created 4 months ago

chore: add extra body & header to openai and anthropic

Change summary

internal/config/config.go          |  2 ++
internal/config/load.go            |  1 +
internal/llm/provider/anthropic.go |  6 ++++++
internal/llm/provider/openai.go    | 10 ++++++----
internal/llm/provider/provider.go  |  2 ++
5 files changed, 17 insertions(+), 4 deletions(-)

Detailed changes

internal/config/config.go

@@ -78,6 +78,8 @@ type ProviderConfig struct {
 
 	// Extra headers to send with each request to the provider.
 	ExtraHeaders map[string]string `json:"extra_headers,omitempty"`
+	// Extra body
+	ExtraBody map[string]any `json:"extra_body,omitempty"`
 
 	// Used to pass extra parameters to the provider.
 	ExtraParams map[string]string `json:"-"`

internal/config/load.go

@@ -171,6 +171,7 @@ func (c *Config) configureProviders(env env.Env, resolver VariableResolver, know
 			Type:         p.Type,
 			Disable:      config.Disable,
 			ExtraHeaders: config.ExtraHeaders,
+			ExtraBody:    config.ExtraBody,
 			ExtraParams:  make(map[string]string),
 			Models:       p.Models,
 		}

internal/llm/provider/anthropic.go

@@ -45,6 +45,12 @@ func createAnthropicClient(opts providerClientOptions, useBedrock bool) anthropi
 	if useBedrock {
 		anthropicClientOptions = append(anthropicClientOptions, bedrock.WithLoadDefaultConfig(context.Background()))
 	}
+	for _, header := range opts.extraHeaders {
+		anthropicClientOptions = append(anthropicClientOptions, option.WithHeaderAdd(header, opts.extraHeaders[header]))
+	}
+	for key, value := range opts.extraBody {
+		anthropicClientOptions = append(anthropicClientOptions, option.WithJSONSet(key, value))
+	}
 	return anthropic.NewClient(anthropicClientOptions...)
 }
 

internal/llm/provider/openai.go

@@ -44,10 +44,12 @@ func createOpenAIClient(opts providerClientOptions) openai.Client {
 		}
 	}
 
-	if opts.extraHeaders != nil {
-		for key, value := range opts.extraHeaders {
-			openaiClientOptions = append(openaiClientOptions, option.WithHeader(key, value))
-		}
+	for key, value := range opts.extraHeaders {
+		openaiClientOptions = append(openaiClientOptions, option.WithHeader(key, value))
+	}
+
+	for extraKey, extraValue := range opts.extraBody {
+		openaiClientOptions = append(openaiClientOptions, option.WithJSONSet(extraKey, extraValue))
 	}
 
 	return openai.NewClient(openaiClientOptions...)

internal/llm/provider/provider.go

@@ -70,6 +70,7 @@ type providerClientOptions struct {
 	systemMessage string
 	maxTokens     int64
 	extraHeaders  map[string]string
+	extraBody     map[string]any
 	extraParams   map[string]string
 }
 
@@ -147,6 +148,7 @@ func NewProvider(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provi
 		config:       cfg,
 		apiKey:       resolvedAPIKey,
 		extraHeaders: cfg.ExtraHeaders,
+		extraBody:    cfg.ExtraBody,
 		model: func(tp config.SelectedModelType) provider.Model {
 			return *config.Get().GetModelByType(tp)
 		},