fix(provider): implement proper provider-per-candidate fallback architecture
This commit fixes a critical bug in the model fallback chain where the `provider` parameter was being completely ignored, causing all fallback candidates to use the same provider instance configured for the primary model. This led to requests being sent to the wrong API endpoints. ## Root Cause The fallback callback at loop.go:522-527 was using `agent.Provider.Chat()` for ALL candidates, ignoring the `provider` parameter passed by the fallback chain. For example: - Primary: cerebras/gpt-oss-120b → agent.Provider configured for api.cerebras.ai - Fallback: ollama/llama3.2 → STILL uses cerebras provider → wrong endpoint! Additionally, the `prompt_cache_key` parameter (an OpenAI-only feature) was being sent to all providers, causing 422 errors from providers like Cerebras that don't support this parameter. ## Changes ### Core Architecture Fix 1. **ProviderRegistry** (pkg/providers/registry.go): New registry that creates and caches provider instances per provider name (e.g., "cerebras", "ollama") 2. **AgentInstance** (pkg/agent/instance.go): Added ProviderRegistry field 3. **Fallback callback** (pkg/agent/loop.go:522-544): Now gets correct provider per candidate via agent.ProviderRegistry.GetProvider(providerName) ### Prompt Cache Fix 1. **Agent loop** (pkg/agent/loop.go): Removed `prompt_cache_key` from generic options map (lines 526, 543, 1009, 1060) 2. **OpenAI compat provider** (pkg/providers/openai_compat/provider.go:158-162): Changed from blacklist (exclude Gemini) to whitelist (only api.openai.com) ### Test Updates - Updated tests to pass new ProviderRegistry parameter ## Impact - Fallback chains with mixed providers (cerebras + ollama + openai) now work - Providers like Cerebras no longer receive unsupported `prompt_cache_key` - Each fallback candidate uses its correctly configured provider Fixes: Model fallback chain sending requests to wrong API endpoints Related: 422 errors from Cerebras due to unsupported prompt_cache_key
This commit is contained in:
parent
5f6a6b5829
commit
af8aa29e05
7 changed files with 174 additions and 55 deletions
|
|
@ -24,7 +24,8 @@ type AgentInstance struct {
|
|||
MaxTokens int
|
||||
Temperature float64
|
||||
ContextWindow int
|
||||
Provider providers.LLMProvider
|
||||
Provider providers.LLMProvider // Default provider for backward compatibility
|
||||
ProviderRegistry *providers.ProviderRegistry // Registry for multi-provider fallback
|
||||
Sessions *session.SessionManager
|
||||
ContextBuilder *ContextBuilder
|
||||
Tools *tools.ToolRegistry
|
||||
|
|
@ -39,6 +40,7 @@ func NewAgentInstance(
|
|||
defaults *config.AgentDefaults,
|
||||
cfg *config.Config,
|
||||
provider providers.LLMProvider,
|
||||
providerRegistry *providers.ProviderRegistry,
|
||||
) *AgentInstance {
|
||||
workspace := resolveAgentWorkspace(agentCfg, defaults)
|
||||
os.MkdirAll(workspace, 0o755)
|
||||
|
|
@ -109,6 +111,7 @@ func NewAgentInstance(
|
|||
Temperature: temperature,
|
||||
ContextWindow: maxTokens,
|
||||
Provider: provider,
|
||||
ProviderRegistry: providerRegistry,
|
||||
Sessions: sessionsManager,
|
||||
ContextBuilder: contextBuilder,
|
||||
Tools: toolsRegistry,
|
||||
|
|
|
|||
|
|
@ -133,7 +133,7 @@ func TestNewAgentInstance_ModelResolution(t *testing.T) {
|
|||
// Use the existing mockProvider from mock_provider_test.go
|
||||
provider := &mockProvider{}
|
||||
|
||||
instance := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
|
||||
instance := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
|
||||
|
||||
require.NotNil(t, instance)
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) {
|
|||
cfg.Agents.Defaults.Temperature = &configuredTemp
|
||||
|
||||
provider := &mockProvider{}
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
|
||||
|
||||
if agent.MaxTokens != 1234 {
|
||||
t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234)
|
||||
|
|
@ -61,7 +61,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenZero(t *testing.T) {
|
|||
cfg.Agents.Defaults.Temperature = &configuredTemp
|
||||
|
||||
provider := &mockProvider{}
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
|
||||
|
||||
if agent.Temperature != 0.0 {
|
||||
t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.0)
|
||||
|
|
@ -87,7 +87,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) {
|
|||
}
|
||||
|
||||
provider := &mockProvider{}
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
|
||||
|
||||
if agent.Temperature != 0.7 {
|
||||
t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.7)
|
||||
|
|
|
|||
|
|
@ -519,11 +519,26 @@ func (al *AgentLoop) runLLMIteration(
|
|||
callLLM := func() (*providers.LLMResponse, error) {
|
||||
if len(agent.Candidates) > 1 && al.fallback != nil {
|
||||
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
|
||||
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
|
||||
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
|
||||
func(ctx context.Context, providerName, model string) (*providers.LLMResponse, error) {
|
||||
// Get the correct provider for this candidate
|
||||
// This is critical: different candidates may use different providers
|
||||
// (e.g., cerebras, ollama, openai), and each requires its own
|
||||
// provider instance with the correct API endpoint configuration.
|
||||
var provider providers.LLMProvider
|
||||
if agent.ProviderRegistry != nil {
|
||||
p, err := agent.ProviderRegistry.GetProvider(providerName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get provider for %q: %w", providerName, err)
|
||||
}
|
||||
provider = p
|
||||
} else {
|
||||
// Fallback to default provider if registry not available
|
||||
// (this shouldn't happen in normal operation)
|
||||
provider = agent.Provider
|
||||
}
|
||||
return provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": agent.Temperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
})
|
||||
},
|
||||
)
|
||||
|
|
@ -540,7 +555,6 @@ func (al *AgentLoop) runLLMIteration(
|
|||
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": agent.Temperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -1006,7 +1020,6 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
|
|||
map[string]any{
|
||||
"max_tokens": 1024,
|
||||
"temperature": 0.3,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
if err == nil {
|
||||
|
|
@ -1057,7 +1070,6 @@ func (al *AgentLoop) summarizeBatch(
|
|||
map[string]any{
|
||||
"max_tokens": 1024,
|
||||
"temperature": 0.3,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -26,20 +26,23 @@ func NewAgentRegistry(
|
|||
resolver: routing.NewRouteResolver(cfg),
|
||||
}
|
||||
|
||||
// Create provider registry from config for proper fallback support
|
||||
providerRegistry := providers.NewProviderRegistry(cfg)
|
||||
|
||||
agentConfigs := cfg.Agents.List
|
||||
if len(agentConfigs) == 0 {
|
||||
implicitAgent := &config.AgentConfig{
|
||||
ID: "main",
|
||||
Default: true,
|
||||
}
|
||||
instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider)
|
||||
instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider, providerRegistry)
|
||||
registry.agents["main"] = instance
|
||||
logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil)
|
||||
} else {
|
||||
for i := range agentConfigs {
|
||||
ac := &agentConfigs[i]
|
||||
id := routing.NormalizeAgentID(ac.ID)
|
||||
instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider)
|
||||
instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider, providerRegistry)
|
||||
registry.agents[id] = instance
|
||||
logger.InfoCF("agent", "Registered agent",
|
||||
map[string]any{
|
||||
|
|
|
|||
|
|
@ -153,10 +153,11 @@ func (p *Provider) Chat(
|
|||
// with the same key and reuse prefix KV cache across calls.
|
||||
// The key is typically the agent ID — stable per agent, shared across requests.
|
||||
// See: https://platform.openai.com/docs/guides/prompt-caching
|
||||
// Prompt caching is only supported by OpenAI-native endpoints.
|
||||
// Gemini and other providers reject unknown fields, so skip for non-OpenAI APIs.
|
||||
// IMPORTANT: Prompt caching is ONLY supported by OpenAI-native endpoints (api.openai.com).
|
||||
// Gemini, Cerebras, and other providers reject unknown fields, so we must skip them.
|
||||
if cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" {
|
||||
if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") {
|
||||
// Only send prompt_cache_key to actual OpenAI endpoints
|
||||
if strings.Contains(p.apiBase, "api.openai.com") {
|
||||
requestBody["prompt_cache_key"] = cacheKey
|
||||
}
|
||||
}
|
||||
|
|
|
|||
100
pkg/providers/registry.go
Normal file
100
pkg/providers/registry.go
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package providers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// ProviderRegistry manages provider instances for different provider names.
|
||||
// It lazily creates providers on-demand and caches them for reuse.
|
||||
// This is essential for proper fallback behavior where different candidates
|
||||
// may use different providers (e.g., cerebras, ollama, openai).
|
||||
type ProviderRegistry struct {
|
||||
cfg *config.Config
|
||||
modelList []config.ModelConfig
|
||||
providers map[string]LLMProvider
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewProviderRegistry creates a new provider registry with the given config.
|
||||
func NewProviderRegistry(cfg *config.Config) *ProviderRegistry {
|
||||
var modelList []config.ModelConfig
|
||||
if cfg != nil {
|
||||
modelList = cfg.ModelList
|
||||
}
|
||||
return &ProviderRegistry{
|
||||
cfg: cfg,
|
||||
modelList: modelList,
|
||||
providers: make(map[string]LLMProvider),
|
||||
}
|
||||
}
|
||||
|
||||
// GetProvider returns a provider for the given provider name (e.g., "openai", "cerebras", "ollama").
|
||||
// If the provider has already been created, it returns the cached instance.
|
||||
// Otherwise, it creates a new provider from the model list config.
|
||||
func (pr *ProviderRegistry) GetProvider(providerName string) (LLMProvider, error) {
|
||||
// Normalize provider name (lowercase)
|
||||
providerName = strings.ToLower(providerName)
|
||||
|
||||
// Check cache first
|
||||
pr.mu.RLock()
|
||||
if provider, ok := pr.providers[providerName]; ok {
|
||||
pr.mu.RUnlock()
|
||||
return provider, nil
|
||||
}
|
||||
pr.mu.RUnlock()
|
||||
|
||||
// Find the model config for this provider
|
||||
var modelCfg *config.ModelConfig
|
||||
for i := range pr.modelList {
|
||||
cfg := &pr.modelList[i]
|
||||
// Extract provider from model string (e.g., "cerebras/gpt-oss-120b" -> "cerebras")
|
||||
protocol, _ := ExtractProtocol(cfg.Model)
|
||||
if strings.EqualFold(protocol, providerName) {
|
||||
modelCfg = cfg
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// If not found in model list, try creating with protocol-only config
|
||||
if modelCfg == nil {
|
||||
modelCfg = &config.ModelConfig{
|
||||
Model: providerName + "/dummy", // Protocol is what matters
|
||||
}
|
||||
}
|
||||
|
||||
// Create the provider
|
||||
provider, _, err := CreateProviderFromConfig(modelCfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create provider for %q: %w", providerName, err)
|
||||
}
|
||||
|
||||
// Cache the provider
|
||||
pr.mu.Lock()
|
||||
pr.providers[providerName] = provider
|
||||
pr.mu.Unlock()
|
||||
|
||||
return provider, nil
|
||||
}
|
||||
|
||||
// GetDefaultProvider returns the provider for the default/primary protocol.
|
||||
// This is used for backward compatibility with code that expects a single provider.
|
||||
func (pr *ProviderRegistry) GetDefaultProvider() (LLMProvider, error) {
|
||||
if len(pr.modelList) == 0 {
|
||||
// No model list configured, return error
|
||||
// Caller should handle this by using a default provider
|
||||
return nil, fmt.Errorf("no model list configured")
|
||||
}
|
||||
|
||||
// Use the first model's provider as default
|
||||
protocol, _ := ExtractProtocol(pr.modelList[0].Model)
|
||||
return pr.GetProvider(protocol)
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue