feat: add userAgent config for ModelConfig

This commit is contained in:
Cytown 2026-04-01 13:02:53 +08:00
parent ff90a65814
commit afd6288124
8 changed files with 66 additions and 29 deletions

View file

@ -598,6 +598,8 @@ type ModelConfig struct {
// existing configs, the field is inferred during load: models with API keys // existing configs, the field is inferred during load: models with API keys
// or the reserved "local-model" name are auto-enabled. // or the reserved "local-model" name are auto-enabled.
Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"` Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`
// UserAgent is the user agent string to use for HTTP requests.
UserAgent string `json:"user_agent,omitempty" yaml:"-"`
// isVirtual marks this model as a virtual model generated from multi-key expansion. // isVirtual marks this model as a virtual model generated from multi-key expansion.
// Virtual models should not be persisted to config files. // Virtual models should not be persisted to config files.

View file

@ -41,15 +41,16 @@ type Provider struct {
apiKey string apiKey string
apiBase string apiBase string
httpClient *http.Client httpClient *http.Client
userAgent string
} }
// NewProvider creates a new Anthropic Messages API provider. // NewProvider creates a new Anthropic Messages API provider.
func NewProvider(apiKey, apiBase string) *Provider { func NewProvider(apiKey, apiBase, userAgent string) *Provider {
return NewProviderWithTimeout(apiKey, apiBase, 0) return NewProviderWithTimeout(apiKey, apiBase, userAgent, 0)
} }
// NewProviderWithTimeout creates a provider with custom request timeout. // NewProviderWithTimeout creates a provider with custom request timeout.
func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provider { func NewProviderWithTimeout(apiKey, apiBase, userAgent string, timeoutSeconds int) *Provider {
baseURL := normalizeBaseURL(apiBase) baseURL := normalizeBaseURL(apiBase)
timeout := defaultRequestTimeout timeout := defaultRequestTimeout
if timeoutSeconds > 0 { if timeoutSeconds > 0 {
@ -57,8 +58,9 @@ func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provide
} }
return &Provider{ return &Provider{
apiKey: apiKey, apiKey: apiKey,
apiBase: baseURL, apiBase: baseURL,
userAgent: userAgent,
httpClient: &http.Client{ httpClient: &http.Client{
Timeout: timeout, Timeout: timeout,
}, },
@ -105,6 +107,9 @@ func (p *Provider) Chat(
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-API-Key", p.apiKey) //nolint:canonicalheader // Anthropic API requires exact header name req.Header.Set("X-API-Key", p.apiKey) //nolint:canonicalheader // Anthropic API requires exact header name
req.Header.Set("Anthropic-Version", defaultAPIVersion) req.Header.Set("Anthropic-Version", defaultAPIVersion)
if p.userAgent != "" {
req.Header.Set("User-Agent", p.userAgent)
}
// Execute request // Execute request
resp, err := p.httpClient.Do(req) resp, err := p.httpClient.Do(req)

View file

@ -411,7 +411,7 @@ func TestNormalizeBaseURL(t *testing.T) {
} }
func TestNewProvider(t *testing.T) { func TestNewProvider(t *testing.T) {
provider := NewProvider("test-key", "https://api.example.com") provider := NewProvider("test-key", "https://api.example.com", "")
if provider == nil { if provider == nil {
t.Fatal("NewProvider() returned nil") t.Fatal("NewProvider() returned nil")
} }
@ -424,7 +424,7 @@ func TestNewProvider(t *testing.T) {
} }
func TestGetDefaultModel(t *testing.T) { func TestGetDefaultModel(t *testing.T) {
provider := NewProvider("test-key", "") provider := NewProvider("test-key", "", "")
got := provider.GetDefaultModel() got := provider.GetDefaultModel()
expected := "claude-sonnet-4.6" expected := "claude-sonnet-4.6"
if got != expected { if got != expected {
@ -743,7 +743,7 @@ func TestProviderChatErrors(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
// Create provider using constructor to ensure proper initialization // Create provider using constructor to ensure proper initialization
provider := NewProvider(tt.apiKey, "https://api.example.com") provider := NewProvider(tt.apiKey, "https://api.example.com", "")
_, err := provider.Chat(context.Background(), tt.messages, nil, "test-model", nil) _, err := provider.Chat(context.Background(), tt.messages, nil, "test-model", nil)
if err == nil { if err == nil {

View file

@ -36,6 +36,7 @@ type Provider struct {
apiKey string apiKey string
apiBase string apiBase string
httpClient *http.Client httpClient *http.Client
userAgent string
} }
// Option configures the Azure Provider. // Option configures the Azure Provider.
@ -50,11 +51,19 @@ func WithRequestTimeout(timeout time.Duration) Option {
} }
} }
// WithUserAgent sets the User-Agent header for requests.
func WithUserAgent(userAgent string) Option {
return func(p *Provider) {
p.userAgent = userAgent
}
}
// NewProvider creates a new Azure OpenAI provider. // NewProvider creates a new Azure OpenAI provider.
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { func NewProvider(apiKey, apiBase, proxy, userAgent string, opts ...Option) *Provider {
p := &Provider{ p := &Provider{
apiKey: apiKey, apiKey: apiKey,
apiBase: strings.TrimRight(apiBase, "/"), apiBase: strings.TrimRight(apiBase, "/"),
userAgent: userAgent,
httpClient: common.NewHTTPClient(proxy), httpClient: common.NewHTTPClient(proxy),
} }
@ -68,9 +77,9 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
} }
// NewProviderWithTimeout creates a new Azure OpenAI provider with a custom request timeout in seconds. // NewProviderWithTimeout creates a new Azure OpenAI provider with a custom request timeout in seconds.
func NewProviderWithTimeout(apiKey, apiBase, proxy string, requestTimeoutSeconds int) *Provider { func NewProviderWithTimeout(apiKey, apiBase, proxy, userAgent string, requestTimeoutSeconds int) *Provider {
return NewProvider( return NewProvider(
apiKey, apiBase, proxy, apiKey, apiBase, proxy, userAgent,
WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second), WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
) )
} }
@ -141,6 +150,9 @@ func (p *Provider) Chat(
if p.apiKey != "" { if p.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("Authorization", "Bearer "+p.apiKey)
} }
if p.userAgent != "" {
req.Header.Set("User-Agent", p.userAgent)
}
resp, err := p.httpClient.Do(req) resp, err := p.httpClient.Do(req)
if err != nil { if err != nil {

View file

@ -46,7 +46,7 @@ func TestProviderChat_AzureURLConstruction(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-gpt5-deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-gpt5-deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -69,7 +69,7 @@ func TestProviderChat_AzureAuthHeader(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-azure-key", server.URL, "") p := NewProvider("test-azure-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -92,7 +92,7 @@ func TestProviderChat_AzureRequestBodyContainsModel(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -112,7 +112,7 @@ func TestProviderChat_AzureUsesMaxOutputTokens(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat( _, err := p.Chat(
t.Context(), t.Context(),
[]Message{{Role: "user", Content: "hi"}}, []Message{{Role: "user", Content: "hi"}},
@ -144,7 +144,7 @@ func TestProviderChat_AzureStoreIsFalse(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -161,7 +161,7 @@ func TestProviderChat_AzureHTTPError(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("bad-key", server.URL, "") p := NewProvider("bad-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err == nil { if err == nil {
t.Fatal("expected error, got nil") t.Fatal("expected error, got nil")
@ -176,7 +176,7 @@ func TestProviderChat_AzureRateLimitError(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err == nil { if err == nil {
t.Fatal("expected error for 429, got nil") t.Fatal("expected error for 429, got nil")
@ -194,7 +194,7 @@ func TestProviderChat_AzureServerError(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err == nil { if err == nil {
t.Fatal("expected error for 500, got nil") t.Fatal("expected error for 500, got nil")
@ -229,7 +229,7 @@ func TestProviderChat_AzureParseTextOutput(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -270,7 +270,7 @@ func TestProviderChat_AzureParseToolCalls(t *testing.T) {
})) }))
defer server.Close() defer server.Close()
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "deployment", nil) out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "deployment", nil)
if err != nil { if err != nil {
t.Fatalf("Chat() error = %v", err) t.Fatalf("Chat() error = %v", err)
@ -287,7 +287,7 @@ func TestProviderChat_AzureParseToolCalls(t *testing.T) {
} }
func TestProvider_AzureEmptyAPIBase(t *testing.T) { func TestProvider_AzureEmptyAPIBase(t *testing.T) {
p := NewProvider("test-key", "", "") p := NewProvider("test-key", "", "", "")
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
if err == nil { if err == nil {
t.Fatal("expected error for empty API base") t.Fatal("expected error for empty API base")
@ -295,21 +295,21 @@ func TestProvider_AzureEmptyAPIBase(t *testing.T) {
} }
func TestProvider_AzureRequestTimeoutDefault(t *testing.T) { func TestProvider_AzureRequestTimeoutDefault(t *testing.T) {
p := NewProvider("test-key", "https://example.com", "") p := NewProvider("test-key", "https://example.com", "", "")
if p.httpClient.Timeout != defaultRequestTimeout { if p.httpClient.Timeout != defaultRequestTimeout {
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout) t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout)
} }
} }
func TestProvider_AzureRequestTimeoutOverride(t *testing.T) { func TestProvider_AzureRequestTimeoutOverride(t *testing.T) {
p := NewProvider("test-key", "https://example.com", "", WithRequestTimeout(300*time.Second)) p := NewProvider("test-key", "https://example.com", "", "", WithRequestTimeout(300*time.Second))
if p.httpClient.Timeout != 300*time.Second { if p.httpClient.Timeout != 300*time.Second {
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 300*time.Second) t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 300*time.Second)
} }
} }
func TestProvider_AzureNewProviderWithTimeout(t *testing.T) { func TestProvider_AzureNewProviderWithTimeout(t *testing.T) {
p := NewProviderWithTimeout("test-key", "https://example.com", "", 180) p := NewProviderWithTimeout("test-key", "https://example.com", "", "", 180)
if p.httpClient.Timeout != 180*time.Second { if p.httpClient.Timeout != 180*time.Second {
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 180*time.Second) t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 180*time.Second)
} }
@ -343,7 +343,7 @@ func TestProviderChat_AzureNativeWebSearchInjection(t *testing.T) {
}, },
} }
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
// With native_search=true: user-defined web_search should be replaced by built-in // With native_search=true: user-defined web_search should be replaced by built-in
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment", _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment",
@ -393,7 +393,7 @@ func TestProviderChat_AzureNoNativeWebSearch(t *testing.T) {
}, },
} }
p := NewProvider("test-key", server.URL, "") p := NewProvider("test-key", server.URL, "", "")
// Without native_search: user-defined web_search should be kept as-is // Without native_search: user-defined web_search should be kept as-is
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment", nil) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment", nil)

View file

@ -138,6 +138,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
apiBase, apiBase,
cfg.Proxy, cfg.Proxy,
cfg.MaxTokensField, cfg.MaxTokensField,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
cfg.ExtraBody, cfg.ExtraBody,
), modelID, nil ), modelID, nil
@ -157,6 +158,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
cfg.APIKey(), cfg.APIKey(),
cfg.APIBase, cfg.APIBase,
cfg.Proxy, cfg.Proxy,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
), modelID, nil ), modelID, nil
@ -214,6 +216,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
apiBase, apiBase,
cfg.Proxy, cfg.Proxy,
cfg.MaxTokensField, cfg.MaxTokensField,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
cfg.ExtraBody, cfg.ExtraBody,
), modelID, nil ), modelID, nil
@ -239,6 +242,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
apiBase, apiBase,
cfg.Proxy, cfg.Proxy,
cfg.MaxTokensField, cfg.MaxTokensField,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
extraBody, extraBody,
), modelID, nil ), modelID, nil
@ -265,6 +269,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
apiBase, apiBase,
cfg.Proxy, cfg.Proxy,
cfg.MaxTokensField, cfg.MaxTokensField,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
cfg.ExtraBody, cfg.ExtraBody,
), modelID, nil ), modelID, nil
@ -281,6 +286,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
return anthropicmessages.NewProviderWithTimeout( return anthropicmessages.NewProviderWithTimeout(
cfg.APIKey(), cfg.APIKey(),
apiBase, apiBase,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
), modelID, nil ), modelID, nil
@ -296,6 +302,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
return anthropicmessages.NewProviderWithTimeout( return anthropicmessages.NewProviderWithTimeout(
cfg.APIKey(), cfg.APIKey(),
apiBase, apiBase,
cfg.UserAgent,
cfg.RequestTimeout, cfg.RequestTimeout,
), modelID, nil ), modelID, nil

View file

@ -24,11 +24,11 @@ func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider {
} }
func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *HTTPProvider { func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *HTTPProvider {
return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, 0, nil) return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, "", 0, nil)
} }
func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
apiKey, apiBase, proxy, maxTokensField string, apiKey, apiBase, proxy, maxTokensField, userAgent string,
requestTimeoutSeconds int, requestTimeoutSeconds int,
extraBody map[string]any, extraBody map[string]any,
) *HTTPProvider { ) *HTTPProvider {
@ -40,6 +40,7 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
openai_compat.WithMaxTokensField(maxTokensField), openai_compat.WithMaxTokensField(maxTokensField),
openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second), openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
openai_compat.WithExtraBody(extraBody), openai_compat.WithExtraBody(extraBody),
openai_compat.WithUserAgent(userAgent),
), ),
} }
} }

View file

@ -36,6 +36,7 @@ type Provider struct {
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models) maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
httpClient *http.Client httpClient *http.Client
extraBody map[string]any // Additional fields to inject into request body extraBody map[string]any // Additional fields to inject into request body
userAgent string
} }
type Option func(*Provider) type Option func(*Provider)
@ -65,6 +66,12 @@ func WithMaxTokensField(maxTokensField string) Option {
} }
} }
func WithUserAgent(userAgent string) Option {
return func(p *Provider) {
p.userAgent = userAgent
}
}
func WithRequestTimeout(timeout time.Duration) Option { func WithRequestTimeout(timeout time.Duration) Option {
return func(p *Provider) { return func(p *Provider) {
if timeout > 0 { if timeout > 0 {
@ -197,6 +204,9 @@ func (p *Provider) Chat(
} }
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
if p.userAgent != "" {
req.Header.Set("User-Agent", p.userAgent)
}
if p.apiKey != "" { if p.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("Authorization", "Bearer "+p.apiKey)
} }