Allow custom headers on OpenAI providers (to be added to more later)

This commit is contained in:
Matt Low 2024-08-12 17:14:53 +00:00
parent fe838f400f
commit 434fc4672b
3 changed files with 11 additions and 5 deletions

View File

@ -16,6 +16,7 @@ import (
type OpenAIClient struct { type OpenAIClient struct {
APIKey string APIKey string
BaseURL string BaseURL string
Headers map[string]string
} }
type ChatCompletionMessage struct { type ChatCompletionMessage struct {
@ -198,6 +199,9 @@ func (c *OpenAIClient) sendRequest(ctx context.Context, r ChatCompletionRequest)
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+c.APIKey) req.Header.Set("Authorization", "Bearer "+c.APIKey)
for header, val := range c.Headers {
req.Header.Set(header, val)
}
client := &http.Client{} client := &http.Client{}
resp, err := client.Do(req) resp, err := client.Do(req)

View File

@ -36,6 +36,7 @@ type Config struct {
BaseURL string `yaml:"baseUrl,omitempty"` BaseURL string `yaml:"baseUrl,omitempty"`
APIKey string `yaml:"apiKey,omitempty"` APIKey string `yaml:"apiKey,omitempty"`
Models []string `yaml:"models"` Models []string `yaml:"models"`
Headers map[string]string `yaml:"headers"`
} `yaml:"providers"` } `yaml:"providers"`
} }

View File

@ -179,6 +179,7 @@ func (c *Context) GetModelProvider(model string) (string, api.ChatCompletionProv
return model, &openai.OpenAIClient{ return model, &openai.OpenAIClient{
BaseURL: url, BaseURL: url,
APIKey: p.APIKey, APIKey: p.APIKey,
Headers: p.Headers,
}, nil }, nil
default: default:
return "", nil, fmt.Errorf("unknown provider kind: %s", p.Kind) return "", nil, fmt.Errorf("unknown provider kind: %s", p.Kind)