Support sending extra headers and changing auth header

This commit introduces support for adding extra HTTP headers that
will be sent on every request to backends of type "openai" and
"ollama". This is useful for providers implementing the OpenAI
API but requiring extra headers, such as Portkey.

For openai backends, the header used for authorization can also now
be modified via the configuration setting `auth_header`. Azure
OpenAI uses the "api-key" header instead of "Authorization".
Previously, aiac used this header whenever the URL for the backend
was not the default OpenAI API URL. This is clearly wrong as it's
really only true for Azure OpenAI, and wrong for providers like
Portkey. The new setting is now used instead, but if the URL is
an Azure OpenAI URL, aiac will automatically know to use "api-key"
instead of "Authorization".

README.md

@@ -153,13 +153,15 @@ default_backend = "official_openai"   # Default backend when one is not selected
 [backends.official_openai]
 type = "openai"
 api_key = "API KEY"
-default_model = "gpt-4o"     # Default model to use for this backend
+default_model = "gpt-4o"              # Default model to use for this backend
 
 [backends.azure_openai]
 type = "openai"
 url = "https://tenant.openai.azure.com/openai/deployments/test"
 api_key = "API KEY"
-api_version = "2023-05-15"   # Optional
+api_version = "2023-05-15"            # Optional
+auth_header = "api-key"               # Default is "Authorization"
+extra_headers = { X-Header-1 = "one", X-Header-2 = "two" }
 
 [backends.aws_staging]
 type = "bedrock"
@@ -177,6 +179,18 @@ type = "ollama"
 url = "http://localhost:11434/api"     # This is the default
 ```
 
+Notes:
+
+1. Every backend can have a default model (via configuration key `default_model`).
+   If not provided, calls that do not define a model will fail.
+2. Backends of type "openai" can change the header used for authorization by
+   providing the `auth_header` setting. This defaults to "Authorization", but
+   Azure OpenAI uses "api-key" instead. When the header is either "Authorization"
+   or "Proxy-Authorization", the header's value for requests will be "Bearer
+   API_KEY". If it's anything else, it'll simply be "API_KEY".
+3. Backends of type "openai" and "ollama" support adding extra headers to every
+   request issued by aiac, by utilizing the `extra_headers` setting.
+
 ### Usage
 
 Once a configuration file is created, you can start generating code and you only

libaiac/config.go

@@ -60,6 +60,10 @@ type BackendConfig struct {
 	// DefaultModel is the name of the model to use by default when a specific
 	// one is not selected.
 	DefaultModel string `toml:"default_model"`
+
+	// ExtraHeaders allows setting extra HTTP headers whenever aiac sends
+	// requests to the backend. Bedrock backends do not support this setting.
+	ExtraHeaders map[string]string `toml:"extra_headers"`
 }
 
 // LoadConfig loads an aiac configuration file from the provided path, which

libaiac/libaiac.go

@@ -133,14 +133,16 @@ func (aiac *Aiac) loadBackend(ctx context.Context, name string) (
 		backend = bedrock.New(cfg)
 	case BackendOllama:
 		backend = ollama.New(&ollama.Options{
-			URL: backendConf.URL,
+			URL:          backendConf.URL,
+			ExtraHeaders: backendConf.ExtraHeaders,
 		})
 	default:
 		// default to openai
 		backend, err = openai.New(&openai.Options{
-			ApiKey:     backendConf.APIKey,
-			URL:        backendConf.URL,
-			APIVersion: backendConf.APIVersion,
+			ApiKey:       backendConf.APIKey,
+			URL:          backendConf.URL,
+			APIVersion:   backendConf.APIVersion,
+			ExtraHeaders: backendConf.ExtraHeaders,
 		})
 		if err != nil {
 			return nil, defaultModel, err

libaiac/ollama/ollama.go

@@ -24,6 +24,10 @@ type Options struct {
 	// URL is the URL of the API server (including the /api path prefix).
 	// Defaults to DefaultAPIURL.
 	URL string
+
+	// ExtraHeaders are extra HTTP headers to send with every request to the
+	// provider.
+	ExtraHeaders map[string]string
 }
 
 // New creates a new instance of the Ollama struct, with the provided
@@ -66,5 +70,9 @@ func New(opts *Options) *Ollama {
 			)
 		})
 
+	for header, value := range opts.ExtraHeaders {
+		cli.HTTPClient.Header(header, value)
+	}
+
 	return cli
 }

libaiac/openai/openai.go

@@ -20,6 +20,7 @@ type OpenAI struct {
 	*requests.HTTPClient
 	apiKey     string
 	apiVersion string
+	authHeader string
 }
 
 // Options is a struct containing all the parameters accepted by the New
@@ -34,8 +35,15 @@ type Options struct {
 	// APIVersion is the version of the OpenAI API to use. Optional.
 	APIVersion string
 
-	// DefaultModel
-	DefaultModel string
+	// AuthHeader allows modifying the header where the API key is sent. This
+	// defaults to Authorization. If it is "Authorization" or
+	// "Proxy-Authorization", the API key is sent with a "Bearer " prefix. If
+	// it's anything else, the API key is sent alone.
+	AuthHeader string
+
+	// ExtraHeaders are extra HTTP headers to send with every request to the
+	// provider.
+	ExtraHeaders map[string]string
 }
 
 // New creates a new instance of the OpenAI struct, with the provided input
@@ -49,23 +57,41 @@ func New(opts *Options) (*OpenAI, error) {
 		return nil, errors.New("OpenAI backends require an API key")
 	}
 
+	// Trim "Bearer " prefix if user accidentally included it, probably by
+	// copy-pasting from somewhere.
+	opts.ApiKey = strings.TrimPrefix(opts.ApiKey, "Bearer ")
+
 	if opts.URL == "" {
 		opts.URL = OpenAIBackend
 	}
 
-	var authHeaderKey string
-	var authHeaderVal string
+	authHeaderKey := "Authorization"
+	authHeaderVal := fmt.Sprintf("Bearer %s", opts.ApiKey)
+
+	// If user provided a different authorization header, use it, and if that
+	// header is neither "Authorization" nor "Proxy-Authorization", remove the
+	// "Bearer " prefix from its value.
+	if opts.AuthHeader != "" && opts.AuthHeader != authHeaderKey {
+		authHeaderKey = opts.AuthHeader
+		if authHeaderKey != "Proxy-Authorization" {
+			authHeaderVal = opts.ApiKey
+		}
+	}
 
-	if opts.URL == OpenAIBackend {
-		authHeaderKey = "Authorization"
-		authHeaderVal = fmt.Sprintf("Bearer %s", opts.ApiKey)
-	} else {
+	// The above section depends on the user telling us to use a different
+	// header for authorization. Previously, though, we used 'api-key' as the
+	// header if the URL was anything other than the OpenAI URL. This worked for
+	// Azure OpenAI users, but since many more providers now implement the same
+	// API (e.g. Portkey), that check was no longer correct. To maintain
+	// backwards compatibility for Azure OpenAI users, though, we can change the
+	// auth header by ourselves if the URL is *.openai.azure.com
+	if opts.AuthHeader == "" && strings.Contains(opts.URL, ".openai.azure.com") {
 		authHeaderKey = "api-key"
 		authHeaderVal = opts.ApiKey
 	}
 
 	backend := &OpenAI{
-		apiKey:     strings.TrimPrefix(opts.ApiKey, "Bearer "),
+		apiKey:     opts.ApiKey,
 		apiVersion: opts.APIVersion,
 	}
 
@@ -101,5 +127,9 @@ func New(opts *Options) (*OpenAI, error) {
 			)
 		})
 
+	for header, value := range opts.ExtraHeaders {
+		backend.HTTPClient.Header(header, value)
+	}
+
 	return backend, nil
 }