feat: add NVIDIA and Azure AI providers

This commit is contained in:
stevef 2026-04-04 07:22:59 +02:00
parent 84e42d6904
commit c4c4cbb4eb
7 changed files with 148 additions and 39 deletions

View file

@ -1,5 +1,3 @@
version: "2"
linters:
default: all
disable:

View file

@ -273,7 +273,7 @@ test: generate
## fmt: Format Go code
fmt:
@$(GOLANGCI_LINT) fmt
@$(GO) fmt ./...
## lint: Run linters
lint:

View file

@ -295,20 +295,44 @@ func DecodeToolCallArguments(raw json.RawMessage, name string) map[string]any {
// --- HTTP response helpers ---
// SafetyFilterError is returned when a request or response is blocked by
// an LLM provider's content safety filters.
type SafetyFilterError struct {
Message string
}
func (e *SafetyFilterError) Error() string {
return e.Message
}
// HandleErrorResponse reads a non-200 response body and returns an appropriate error.
func HandleErrorResponse(resp *http.Response, apiBase string) error {
contentType := resp.Header.Get("Content-Type")
body, readErr := io.ReadAll(io.LimitReader(resp.Body, 256))
body, readErr := io.ReadAll(io.LimitReader(resp.Body, 1024)) // Increased limit for detailed error bodies
if readErr != nil {
return fmt.Errorf("failed to read response: %w", readErr)
}
if LooksLikeHTML(body, contentType) {
return WrapHTMLResponseError(resp.StatusCode, body, contentType, apiBase)
}
bodyStr := string(body)
bodyLower := strings.ToLower(bodyStr)
// Detect content safety filters (Azure, OpenAI, etc.)
if strings.Contains(bodyLower, "content_filter") ||
strings.Contains(bodyLower, "content management policy") ||
strings.Contains(bodyLower, "safety filter") ||
strings.Contains(bodyLower, "pii filter") {
return &SafetyFilterError{
Message: "request blocked by provider safety filters: " + ResponsePreview(body, 256),
}
}
return fmt.Errorf(
"API request failed:\n Status: %d\n Body: %s",
resp.StatusCode,
ResponsePreview(body, 128),
ResponsePreview(body, 512),
)
}

View file

@ -217,11 +217,12 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
}
return provider, modelID, nil
case "litellm", "lmstudio", "openrouter", "groq", "zhipu", "gemini", "nvidia", "venice",
case "litellm", "lmstudio", "openrouter", "groq", "zhipu", "gemini", "venice",
"ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras",
"vivgrid", "volcengine", "vllm", "qwen", "qwen-intl", "qwen-international", "dashscope-intl",
"qwen-us", "dashscope-us", "mistral", "avian", "longcat", "modelscope", "novita",
"coding-plan", "alibaba-coding", "qwen-coding", "mimo":
// All other OpenAI-compatible HTTP providers
if cfg.APIKey() == "" && cfg.APIBase == "" && !isEmptyAPIKeyAllowed(protocol) {
return nil, "", fmt.Errorf("api_key or api_base is required for HTTP-based protocol %q", protocol)
@ -240,6 +241,38 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
cfg.ExtraBody,
), modelID, nil
case "nvidia":
apiBase := cfg.APIBase
if apiBase == "" {
apiBase = getDefaultAPIBase(protocol)
}
p := NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
cfg.APIKey(),
apiBase,
cfg.Proxy,
cfg.MaxTokensField,
userAgent,
cfg.RequestTimeout,
cfg.ExtraBody,
)
// NVIDIA sometimes prefers api-key header or has issues with Bearer in some environments
p.SetUseAzureHeaders(false) // NVIDIA main gateway prefers standard Bearer headers; api-key causes 404s
return p, "nvidia/" + modelID, nil
case "azure-ai", "azure-foundry":
// Azure AI Foundry / Studio compatible with OpenAI API format,
// but using api-key header instead of Authorization: Bearer.
if cfg.APIKey() == "" && cfg.APIBase == "" {
return nil, "", fmt.Errorf("api_key or api_base is required for protocol %q", protocol)
}
return NewAzureAIProvider(
cfg.APIKey(),
cfg.APIBase,
cfg.Proxy,
userAgent,
cfg.RequestTimeout,
), modelID, nil
case "minimax":
// Minimax requires reasoning_split: true in the request body
if cfg.APIKey() == "" && cfg.APIBase == "" {

View file

@ -17,9 +17,9 @@ type HTTPProvider struct {
delegate *openai_compat.Provider
}
func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider {
func NewHTTPProvider(apiKey, apiBase, proxy, userAgent string) *HTTPProvider {
return &HTTPProvider{
delegate: openai_compat.NewProvider(apiKey, apiBase, proxy),
delegate: openai_compat.NewProvider(apiKey, apiBase, proxy, openai_compat.WithUserAgent(userAgent)),
}
}
@ -45,6 +45,19 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
}
}
func NewAzureAIProvider(apiKey, apiBase, proxy, userAgent string, requestTimeoutSeconds int) *HTTPProvider {
return &HTTPProvider{
delegate: openai_compat.NewProvider(
apiKey,
apiBase,
proxy,
openai_compat.WithAzureHeaders(true),
openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
openai_compat.WithUserAgent(userAgent),
),
}
}
func (p *HTTPProvider) Chat(
ctx context.Context,
messages []Message,
@ -72,6 +85,10 @@ func (p *HTTPProvider) GetDefaultModel() string {
return ""
}
func (p *HTTPProvider) SetUseAzureHeaders(use bool) {
p.delegate.SetUseAzureHeaders(use)
}
func (p *HTTPProvider) SupportsNativeSearch() bool {
return p.delegate.SupportsNativeSearch()
}

View file

@ -11,6 +11,7 @@ import (
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/providers/common"
@ -31,12 +32,14 @@ type (
)
type Provider struct {
apiKey string
apiBase string
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
httpClient *http.Client
extraBody map[string]any // Additional fields to inject into request body
userAgent string
apiKey string
apiBase string
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
httpClient *http.Client
extraBody map[string]any // Additional fields to inject into request body
userAgent string
useAzureHeaders bool // Use api-key header instead of Authorization: Bearer
mu sync.RWMutex // Protect useAzureHeaders
}
type Option func(*Provider)
@ -44,21 +47,23 @@ type Option func(*Provider)
const defaultRequestTimeout = common.DefaultRequestTimeout
var stripModelPrefixProviders = map[string]struct{}{
"litellm": {},
"venice": {},
"moonshot": {},
"nvidia": {},
"groq": {},
"ollama": {},
"deepseek": {},
"google": {},
"openrouter": {},
"zhipu": {},
"mistral": {},
"vivgrid": {},
"minimax": {},
"novita": {},
"lmstudio": {},
"litellm": {},
"venice": {},
"moonshot": {},
"nvidia": {},
"groq": {},
"ollama": {},
"deepseek": {},
"google": {},
"openrouter": {},
"zhipu": {},
"mistral": {},
"vivgrid": {},
"minimax": {},
"novita": {},
"lmstudio": {},
"azure-ai": {},
"azure-foundry": {},
}
func WithMaxTokensField(maxTokensField string) Option {
@ -87,6 +92,18 @@ func WithExtraBody(extraBody map[string]any) Option {
}
}
func WithAzureHeaders(use bool) Option {
return func(p *Provider) {
p.useAzureHeaders = use
}
}
func (p *Provider) SetUseAzureHeaders(use bool) {
p.mu.Lock()
defer p.mu.Unlock()
p.useAzureHeaders = use
}
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
p := &Provider{
apiKey: apiKey,
@ -209,7 +226,11 @@ func (p *Provider) Chat(
req.Header.Set("User-Agent", p.userAgent)
}
if p.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+p.apiKey)
if p.useAzureHeaders {
req.Header.Set("api-key", p.apiKey)
} else {
req.Header.Set("Authorization", "Bearer "+p.apiKey)
}
}
resp, err := p.httpClient.Do(req)
@ -255,7 +276,11 @@ func (p *Provider) ChatStream(
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
if p.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+p.apiKey)
if p.useAzureHeaders {
req.Header.Set("api-key", p.apiKey)
} else {
req.Header.Set("Authorization", "Bearer "+p.apiKey)
}
}
// Use a client without Timeout for streaming — the http.Client.Timeout covers
@ -415,12 +440,22 @@ func parseStreamResponse(
}
func normalizeModel(model, apiBase string) string {
before, after, ok := strings.Cut(model, "/")
if !ok {
if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") {
return model
}
if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") {
// NVIDIA endpoints (integrate.api.nvidia.com) require the provider prefix
// (e.g., nvidia/, meta/, mistral/) for routing. Do not strip them.
// We also re-add the prefix if it was likely stripped by the agent's protocol resolution logic.
if strings.Contains(strings.ToLower(apiBase), ".nvidia.com") {
if !strings.Contains(model, "/") {
return "nvidia/" + model
}
return model
}
before, after, ok := strings.Cut(model, "/")
if !ok {
return model
}
@ -456,7 +491,7 @@ func isNativeSearchHost(apiBase string) bool {
return false
}
host := u.Hostname()
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
return host == "api.openai.com"
}
// supportsPromptCacheKey reports whether the given API base is known to
@ -469,5 +504,7 @@ func supportsPromptCacheKey(apiBase string) bool {
return false
}
host := u.Hostname()
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
// Strictly limit to OpenAI official. Azure OpenAI often rejects this field
// depending on model version and region, causing 400 errors.
return host == "api.openai.com"
}

View file

@ -923,8 +923,8 @@ func TestSupportsPromptCacheKey(t *testing.T) {
}{
{"https://api.openai.com/v1", true},
{"https://api.openai.com/v1/", true},
{"https://myresource.openai.azure.com/openai/deployments/gpt-4", true},
{"https://eastus.openai.azure.com/v1", true},
{"https://myresource.openai.azure.com/openai/deployments/gpt-4", false},
{"https://eastus.openai.azure.com/v1", false},
{"https://api.mistral.ai/v1", false},
{"https://generativelanguage.googleapis.com/v1beta", false},
{"https://api.deepseek.com/v1", false},
@ -995,7 +995,7 @@ func TestIsNativeSearchHost(t *testing.T) {
want bool
}{
{"https://api.openai.com/v1", true},
{"https://myresource.openai.azure.com/openai/deployments/gpt-4", true},
{"https://myresource.openai.azure.com/openai/deployments/gpt-4", false},
{"https://api.mistral.ai/v1", false},
{"https://api.deepseek.com/v1", false},
{"https://api.groq.com/openai/v1", false},