feat: add NVIDIA and Azure AI providers
This commit is contained in:
parent
84e42d6904
commit
c4c4cbb4eb
7 changed files with 148 additions and 39 deletions
|
|
@ -1,5 +1,3 @@
|
|||
version: "2"
|
||||
|
||||
linters:
|
||||
default: all
|
||||
disable:
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -273,7 +273,7 @@ test: generate
|
|||
|
||||
## fmt: Format Go code
|
||||
fmt:
|
||||
@$(GOLANGCI_LINT) fmt
|
||||
@$(GO) fmt ./...
|
||||
|
||||
## lint: Run linters
|
||||
lint:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 == "" {
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue