mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: fetch providers over DRPC (#26650)
Closes [AIGOV-455](https://linear.app/codercom/issue/AIGOV-455/extend-drpc-with-buildproviders). ## Why The AI Gateway (`aibridged`) is being split into a standalone process that must not touch the database. `coderd` stays the source of truth and seeds the `ai_providers` / `ai_provider_keys` tables from the environment. This PR adds a DRPC call so the gateway fetches provider config from `coderd` instead of reading the DB, for both the embedded and standalone daemons. ## What - **Proto:** new `ProviderConfigurator` service with a unary `GetAIProviders` RPC, plus `AIProvider` / `AIProviderBedrock` messages. `CurrentMinor` bumped to 1 (additive). - **Server (`coderd/aibridgedserver`):** `GetAIProviders` runs a read-only `InTx` under `LockIDAIProvidersEnvSeed` so it never returns a mid-seed snapshot, reads providers (incl. disabled) plus keys for enabled ones, and maps to proto under `dbauthz.AsAIBridged`. Unmappable rows are skipped and logged; plaintext keys and Bedrock secrets are never logged. - **Client:** `DRPCProviderConfiguratorClient` wired into the client union, `dialer.go`, and `CreateInMemoryAIBridgeServer`. - **cli:** `BuildProvidersFromProto` maps the response through the existing DB-neutral `buildProvider`. A shared `poolRPCReloader` does the fetch/build/replace for both daemons: the embedded daemon reloads on every `ai_providers` change and fails startup if it cannot subscribe; the standalone gateway drives the same reloader once at startup, retrying until success and staying interruptible. - **Dead code removed:** `BuildProvidersFromConfig`, `ProvidersFromConfig`, `AIProviderFromConfig`, and the DB-read `BuildProviders` path.
This commit is contained in:
Danny Kopping
+171
-169
@@ -6,7 +6,6 @@ import (
|
||||
"context"
|
||||
"slices"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
@@ -16,9 +15,8 @@ import (
|
||||
"github.com/coder/coder/v2/aibridge/keypool"
|
||||
"github.com/coder/coder/v2/coderd"
|
||||
"github.com/coder/coder/v2/coderd/aibridged"
|
||||
"github.com/coder/coder/v2/coderd/aibridged/proto"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/tracing"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
@@ -26,11 +24,23 @@ import (
|
||||
)
|
||||
|
||||
// newAIBridgeDaemon constructs the in-memory aibridge daemon and wires
|
||||
// up a subscription that hot-reloads the provider pool from the
|
||||
// database on every ai_providers change event. The returned unsubscribe
|
||||
// up a subscription that hot-reloads the provider pool over the in-memory
|
||||
// RPC on every ai_providers change event. The returned unsubscribe
|
||||
// function tears down the subscription; callers must invoke it
|
||||
// alongside Server.Close on shutdown.
|
||||
func newAIBridgeDaemon(coderAPI *coderd.API, providers []aibridge.Provider, cfg codersdk.AIBridgeConfig, reg prometheus.Registerer, metrics *aibridge.Metrics) (*aibridged.Server, func(), error) {
|
||||
//
|
||||
// Reloads fetch the provider set from coderd over the in-memory DRPC
|
||||
// (GetAIProviders) rather than reading the database directly, so embedded and
|
||||
// standalone gateways construct providers identically. Pubsub remains the
|
||||
// hot-reload trigger.
|
||||
//
|
||||
// SubscribeProviderReload performs a best-effort initial reload synchronously,
|
||||
// so the pool is populated before this returns whenever the fetch succeeds.
|
||||
// That reload blocks on srv.Client(), but the embedded daemon's connection is
|
||||
// an in-memory pipe that comes up immediately, and the env seed (which holds
|
||||
// the seed lock) has already completed earlier in startup, so the wait is
|
||||
// negligible.
|
||||
func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg prometheus.Registerer, metrics *aibridge.Metrics) (*aibridged.Server, func(), error) {
|
||||
ctx := context.Background()
|
||||
coderAPI.Logger.Debug(ctx, "starting in-memory aibridge daemon")
|
||||
|
||||
@@ -39,8 +49,10 @@ func newAIBridgeDaemon(coderAPI *coderd.API, providers []aibridge.Provider, cfg
|
||||
providerMetrics := aibridged.NewMetrics(reg)
|
||||
tracer := coderAPI.TracerProvider.Tracer(tracing.TracerName)
|
||||
|
||||
// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
|
||||
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger.Named("pool"), metrics, tracer) // TODO: configurable size.
|
||||
// Create an empty pool for reusable stateful [aibridge.RequestBridge]
|
||||
// instances (one per user). The reloader populates it via the initial
|
||||
// reload below.
|
||||
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger.Named("pool"), metrics, tracer) // TODO: configurable size.
|
||||
if err != nil {
|
||||
return nil, nil, xerrors.Errorf("create request pool: %w", err)
|
||||
}
|
||||
@@ -48,147 +60,121 @@ func newAIBridgeDaemon(coderAPI *coderd.API, providers []aibridge.Provider, cfg
|
||||
// Report current key pool state per provider at scrape time.
|
||||
reg.MustRegister(keypool.NewStateCollector(pool.KeyPools))
|
||||
|
||||
// Subscribe to ai_providers change events so the pool tracks the
|
||||
// database without a restart. The boot-time `providers` snapshot
|
||||
// derives from env config and serves as a fallback if the database
|
||||
// load fails inside the reloader.
|
||||
reloader := &poolDBReloader{
|
||||
pool: pool,
|
||||
db: coderAPI.Database,
|
||||
cfg: cfg,
|
||||
logger: logger.Named("provider-loader"),
|
||||
aibridgeMetrics: metrics,
|
||||
providerMetrics: providerMetrics,
|
||||
}
|
||||
unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload"))
|
||||
if err != nil {
|
||||
// Pool is still usable with the boot-time snapshot; subscription
|
||||
// failure is logged but not fatal so the daemon still serves.
|
||||
logger.Warn(ctx, "subscribe to ai providers change channel", slog.Error(err))
|
||||
unsubscribe = func() {}
|
||||
}
|
||||
|
||||
// Create daemon.
|
||||
// Create daemon. Construct it before subscribing so the reloader can use
|
||||
// srv.Client() to fetch providers over the in-memory RPC.
|
||||
srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
|
||||
return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
|
||||
}, logger, tracer)
|
||||
if err != nil {
|
||||
unsubscribe()
|
||||
return nil, nil, xerrors.Errorf("start in-memory aibridge daemon: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe to ai_providers change events so the pool tracks the database
|
||||
// without a restart, and perform the initial reload. The reload data path
|
||||
// is the in-memory RPC.
|
||||
reloader := NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("provider-loader"), metrics, providerMetrics)
|
||||
unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload"))
|
||||
if err != nil {
|
||||
// Without the subscription the pool can never track provider changes,
|
||||
// so fail startup rather than serve a permanently stale snapshot.
|
||||
_ = srv.Close()
|
||||
return nil, nil, xerrors.Errorf("subscribe to ai providers change channel: %w", err)
|
||||
}
|
||||
|
||||
return srv, unsubscribe, nil
|
||||
}
|
||||
|
||||
// poolDBReloader implements [aibridged.ProviderReloader] by loading
|
||||
// the live provider set from the database and forwarding it to the
|
||||
// pool.
|
||||
type poolDBReloader struct {
|
||||
// poolRPCReloader implements [aibridged.ProviderReloader] by fetching the
|
||||
// live provider set from coderd over a DRPC client and forwarding it to the
|
||||
// pool. It is shared by the embedded daemon (in-memory RPC, pubsub-triggered)
|
||||
// and the standalone gateway (WebSocket RPC, retried at startup) so the fetch,
|
||||
// build, replace, and reload-metric accounting live in one place.
|
||||
type poolRPCReloader struct {
|
||||
pool *aibridged.CachedBridgePool
|
||||
db database.Store
|
||||
client func() (aibridged.DRPCClient, error)
|
||||
cfg codersdk.AIBridgeConfig
|
||||
logger slog.Logger
|
||||
aibridgeMetrics *aibridge.Metrics
|
||||
providerMetrics *aibridged.Metrics
|
||||
}
|
||||
|
||||
func (r *poolDBReloader) Reload(ctx context.Context) error {
|
||||
// NewPoolRPCReloader builds an [aibridged.ProviderReloader] that fetches the
|
||||
// provider set over the DRPC client returned by client and replaces pool's
|
||||
// providers, recording reload metrics against providerMetrics.
|
||||
func NewPoolRPCReloader(
|
||||
pool *aibridged.CachedBridgePool,
|
||||
client func() (aibridged.DRPCClient, error),
|
||||
cfg codersdk.AIBridgeConfig,
|
||||
logger slog.Logger,
|
||||
aibridgeMetrics *aibridge.Metrics,
|
||||
providerMetrics *aibridged.Metrics,
|
||||
) aibridged.ProviderReloader {
|
||||
return &poolRPCReloader{
|
||||
pool: pool,
|
||||
client: client,
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
aibridgeMetrics: aibridgeMetrics,
|
||||
providerMetrics: providerMetrics,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *poolRPCReloader) Reload(ctx context.Context) error {
|
||||
r.providerMetrics.RecordReloadAttempt()
|
||||
providers, outcomes, err := BuildProviders(ctx, r.db, r.cfg, r.logger, r.aibridgeMetrics)
|
||||
// r.client() blocks until the daemon is connected to coderd.
|
||||
client, err := r.client()
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get ai-gateway client: %w", err)
|
||||
}
|
||||
resp, err := client.GetAIProviders(ctx, &proto.GetAIProvidersRequest{})
|
||||
if err != nil {
|
||||
// Keep the previous snapshot in place: dropping all providers
|
||||
// because the DB read failed would compound the visible failure
|
||||
// mode beyond the operator's actual misconfiguration.
|
||||
return xerrors.Errorf("load ai providers from database: %w", err)
|
||||
// because the fetch failed would compound the visible failure mode
|
||||
// beyond the operator's actual misconfiguration.
|
||||
return xerrors.Errorf("fetch ai providers: %w", err)
|
||||
}
|
||||
providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), r.cfg, r.logger, r.aibridgeMetrics)
|
||||
r.pool.ReplaceProviders(providers)
|
||||
r.providerMetrics.RecordReloadSuccess(outcomes)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BuildProviders loads all ai_providers rows (enabled and disabled),
|
||||
// attaches keys to enabled rows, and constructs the equivalent
|
||||
// [aibridge.Provider] instances. The database is the single source of
|
||||
// truth for runtime provider configuration.
|
||||
// BuildProvidersFromProto constructs the runtime [aibridge.Provider] set from
|
||||
// proto provider configuration.
|
||||
//
|
||||
// Disabled rows produce a Provider stub with Enabled() == false so the
|
||||
// Disabled entries produce a Provider stub with Enabled() == false so the
|
||||
// bridge can answer requests targeting them with a 503 sentinel.
|
||||
//
|
||||
// Per-provider construction errors are logged and the offending row is
|
||||
// excluded from the returned snapshot; only a failure of the DB query
|
||||
// itself is propagated. This keeps a single misconfigured row from
|
||||
// taking the whole daemon down.
|
||||
func BuildProviders(ctx context.Context, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) {
|
||||
//nolint:gocritic // AsAIBridged has a minimal permission set for this purpose.
|
||||
authCtx := dbauthz.AsAIBridged(ctx)
|
||||
|
||||
var rows []database.AIProvider
|
||||
keysByProvider := make(map[uuid.UUID][]database.AIProviderKey)
|
||||
|
||||
// Wrap both queries in a read-only transaction so the provider list
|
||||
// and the key list are consistent with each other.
|
||||
err := db.InTx(func(tx database.Store) error {
|
||||
var err error
|
||||
rows, err = tx.GetAIProviders(authCtx, database.GetAIProvidersParams{
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("load ai providers: %w", err)
|
||||
}
|
||||
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Load keys only for the enabled providers to avoid materializing
|
||||
// secrets for disabled rows.
|
||||
ids := make([]uuid.UUID, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
if !r.Enabled {
|
||||
continue
|
||||
}
|
||||
ids = append(ids, r.ID)
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
keyRows, err := tx.GetAIProviderKeysByProviderIDs(authCtx, ids)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("load ai provider keys: %w", err)
|
||||
}
|
||||
for _, k := range keyRows {
|
||||
keysByProvider[k.ProviderID] = append(keysByProvider[k.ProviderID], k)
|
||||
}
|
||||
return nil
|
||||
}, &database.TxOptions{ReadOnly: true, TxIdentifier: "build_ai_providers"})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
providers := make([]aibridge.Provider, 0, len(rows))
|
||||
outcomes := make([]aibridged.ProviderOutcome, 0, len(rows))
|
||||
// Per-provider construction errors are logged and the offending entry is
|
||||
// excluded from the returned snapshot; this keeps a single misconfigured
|
||||
// provider from taking the whole daemon down. The returned outcomes mirror the
|
||||
// per-provider status for metrics reporting.
|
||||
func BuildProvidersFromProto(ctx context.Context, protoProviders []*proto.AIProvider, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics) ([]aibridge.Provider, []aibridged.ProviderOutcome) {
|
||||
providers := make([]aibridge.Provider, 0, len(protoProviders))
|
||||
outcomes := make([]aibridged.ProviderOutcome, 0, len(protoProviders))
|
||||
enabledCount := 0
|
||||
for _, row := range rows {
|
||||
for _, pp := range protoProviders {
|
||||
spec := protoToProviderSpec(pp)
|
||||
outcome := aibridged.ProviderOutcome{
|
||||
Name: row.Name,
|
||||
Type: string(row.Type),
|
||||
Name: spec.Name,
|
||||
Type: string(spec.Type),
|
||||
}
|
||||
if row.Enabled {
|
||||
if spec.Enabled {
|
||||
enabledCount++
|
||||
}
|
||||
prov, err := buildAIProviderFromRow(ctx, row, keysByProvider[row.ID], cfg, metrics)
|
||||
prov, err := buildProvider(ctx, spec, cfg, metrics)
|
||||
if err != nil {
|
||||
outcome.Status = aibridged.ProviderStatusError
|
||||
outcome.Err = err
|
||||
outcomes = append(outcomes, outcome)
|
||||
logger.Error(ctx, "skipping misconfigured ai provider",
|
||||
slog.F("provider_id", row.ID),
|
||||
slog.F("provider_name", row.Name),
|
||||
slog.F("provider_type", string(row.Type)),
|
||||
slog.F("provider_name", spec.Name),
|
||||
slog.F("provider_type", string(spec.Type)),
|
||||
slog.Error(err),
|
||||
)
|
||||
continue
|
||||
}
|
||||
if row.Enabled {
|
||||
if spec.Enabled {
|
||||
outcome.Status = aibridged.ProviderStatusEnabled
|
||||
} else {
|
||||
outcome.Status = aibridged.ProviderStatusDisabled
|
||||
@@ -201,28 +187,56 @@ func BuildProviders(ctx context.Context, db database.Store, cfg codersdk.AIBridg
|
||||
logger.Warn(ctx, "all enabled ai providers failed to build; only disabled providers remain")
|
||||
}
|
||||
|
||||
return providers, outcomes, nil
|
||||
return providers, outcomes
|
||||
}
|
||||
|
||||
// buildAIProviderFromRow decodes the settings blob and constructs the
|
||||
// appropriate [aibridge.Provider] for a single ai_providers row.
|
||||
// Disabled rows return a Provider stub carrying only Name and
|
||||
// Disabled: true; settings decode, key loading, and credential checks
|
||||
// are skipped because the provider will never call upstream.
|
||||
func buildAIProviderFromRow(
|
||||
ctx context.Context,
|
||||
row database.AIProvider,
|
||||
keys []database.AIProviderKey,
|
||||
cfg codersdk.AIBridgeConfig,
|
||||
metrics *aibridge.Metrics,
|
||||
) (aibridge.Provider, error) {
|
||||
if !row.Enabled {
|
||||
return disabledProviderFromRow(row)
|
||||
// protoToProviderSpec maps a proto [proto.AIProvider] into the database-neutral
|
||||
// [aiProviderSpec] consumed by [buildProvider]. Keys and Bedrock settings are
|
||||
// only meaningful for enabled providers; disabled providers carry neither over
|
||||
// the wire.
|
||||
func protoToProviderSpec(pp *proto.AIProvider) aiProviderSpec {
|
||||
spec := aiProviderSpec{
|
||||
Type: database.AIProviderType(pp.GetType()),
|
||||
Name: pp.GetName(),
|
||||
Enabled: pp.GetEnabled(),
|
||||
BaseURL: pp.GetBaseUrl(),
|
||||
Keys: pp.GetKeys(),
|
||||
}
|
||||
if b := pp.GetBedrock(); b != nil {
|
||||
bedrock := codersdk.NewAIProviderBedrockSettings(
|
||||
b.GetRegion(),
|
||||
b.GetAccessKey(),
|
||||
b.GetAccessKeySecret(),
|
||||
b.GetModel(),
|
||||
b.GetSmallFastModel(),
|
||||
)
|
||||
bedrock.RoleARN = b.GetRoleArn()
|
||||
spec.Bedrock = ptr.Ref(bedrock)
|
||||
}
|
||||
return spec
|
||||
}
|
||||
|
||||
settings, err := db2sdk.AIProviderSettings(row.Settings)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("decode settings: %w", err)
|
||||
// aiProviderSpec is a database-neutral description of a single provider,
|
||||
// carrying exactly the inputs [buildProvider] needs. The RPC path
|
||||
// ([protoToProviderSpec]) maps the proto provider into this shape so the
|
||||
// per-type construction logic stays in one place.
|
||||
type aiProviderSpec struct {
|
||||
Type database.AIProviderType
|
||||
Name string
|
||||
Enabled bool
|
||||
BaseURL string
|
||||
// Keys holds bearer API keys for non-Bedrock providers.
|
||||
Keys []string
|
||||
// Bedrock holds Bedrock-specific settings when the provider targets
|
||||
// AWS Bedrock; nil otherwise.
|
||||
Bedrock *codersdk.AIProviderBedrockSettings
|
||||
}
|
||||
|
||||
// buildProvider constructs the appropriate [aibridge.Provider] for a
|
||||
// single provider spec, independent of where the spec was sourced from.
|
||||
func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBridgeConfig, metrics *aibridge.Metrics) (aibridge.Provider, error) {
|
||||
if !spec.Enabled {
|
||||
return aibridge.NewDisabledProviderStub(spec.Name, string(spec.Type)), nil
|
||||
}
|
||||
|
||||
cbCfg := circuitBreakerConfig(cfg)
|
||||
@@ -235,27 +249,27 @@ func buildAIProviderFromRow(
|
||||
// provider because chatd configures them against their
|
||||
// OpenAI-compatible endpoints. Bedrock routes through the Anthropic
|
||||
// provider with a Bedrock discriminator in Settings.
|
||||
switch row.Type {
|
||||
switch spec.Type {
|
||||
case database.AIProviderTypeOpenai,
|
||||
database.AIProviderTypeAzure,
|
||||
database.AIProviderTypeGoogle,
|
||||
database.AIProviderTypeOpenaiCompat,
|
||||
database.AIProviderTypeOpenrouter,
|
||||
database.AIProviderTypeVercel:
|
||||
if len(keys) == 0 && !cfg.AllowBYOK.Value() {
|
||||
return nil, xerrors.Errorf("%s provider has no api keys configured and BYOK is not enabled", row.Type)
|
||||
if len(spec.Keys) == 0 && !cfg.AllowBYOK.Value() {
|
||||
return nil, xerrors.Errorf("%s provider has no api keys configured and BYOK is not enabled", spec.Type)
|
||||
}
|
||||
var pool *keypool.Pool
|
||||
if len(keys) > 0 {
|
||||
if len(spec.Keys) > 0 {
|
||||
var err error
|
||||
pool, err = buildAIProviderKeyPool(row.Name, keys, metrics)
|
||||
pool, err = buildAIProviderKeyPool(spec.Name, spec.Keys, metrics)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("%s key pool: %w", row.Type, err)
|
||||
return nil, xerrors.Errorf("%s key pool: %w", spec.Type, err)
|
||||
}
|
||||
}
|
||||
return aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
|
||||
Name: row.Name,
|
||||
BaseURL: row.BaseUrl,
|
||||
Name: spec.Name,
|
||||
BaseURL: spec.BaseURL,
|
||||
KeyPool: pool,
|
||||
APIDumpDir: dumpDir,
|
||||
CircuitBreaker: cbCfg,
|
||||
@@ -263,31 +277,31 @@ func buildAIProviderFromRow(
|
||||
}), nil
|
||||
|
||||
case database.AIProviderTypeAnthropic, database.AIProviderTypeBedrock:
|
||||
bedrock := bedrockConfigFromRow(row, settings)
|
||||
// A row typed 'bedrock' authenticates exclusively via settings;
|
||||
bedrock := bedrockConfig(spec.BaseURL, spec.Bedrock)
|
||||
// A spec typed 'bedrock' authenticates exclusively via settings;
|
||||
// without populated Bedrock credentials it cannot make upstream
|
||||
// calls, so refuse rather than falling back to an unsigned
|
||||
// Anthropic client.
|
||||
if row.Type == database.AIProviderTypeBedrock && bedrock == nil {
|
||||
if spec.Type == database.AIProviderTypeBedrock && bedrock == nil {
|
||||
return nil, xerrors.New("bedrock provider has no bedrock credentials configured")
|
||||
}
|
||||
// Bedrock-backed Anthropic authenticates via AWS credentials in
|
||||
// the settings blob, not the api_keys table. A bearer-token
|
||||
// Anthropic without any key cannot make upstream calls.
|
||||
if bedrock == nil && len(keys) == 0 && !cfg.AllowBYOK.Value() {
|
||||
// the settings blob, not bearer keys. A bearer-token Anthropic
|
||||
// without any key cannot make upstream calls.
|
||||
if bedrock == nil && len(spec.Keys) == 0 && !cfg.AllowBYOK.Value() {
|
||||
return nil, xerrors.New("anthropic provider has no api keys, no bedrock credentials, and BYOK is not enabled")
|
||||
}
|
||||
var pool *keypool.Pool
|
||||
if len(keys) > 0 {
|
||||
if len(spec.Keys) > 0 {
|
||||
var err error
|
||||
pool, err = buildAIProviderKeyPool(row.Name, keys, metrics)
|
||||
pool, err = buildAIProviderKeyPool(spec.Name, spec.Keys, metrics)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("anthropic key pool: %w", err)
|
||||
}
|
||||
}
|
||||
return aibridge.NewAnthropicProvider(ctx, aibridge.AnthropicConfig{
|
||||
Name: row.Name,
|
||||
BaseURL: row.BaseUrl,
|
||||
Name: spec.Name,
|
||||
BaseURL: spec.BaseURL,
|
||||
KeyPool: pool,
|
||||
APIDumpDir: dumpDir,
|
||||
CircuitBreaker: cbCfg,
|
||||
@@ -298,52 +312,40 @@ func buildAIProviderFromRow(
|
||||
// Copilot is always BYOK; the per-user token is supplied on each
|
||||
// request via the Authorization header, so no keypool is built.
|
||||
return aibridge.NewCopilotProvider(aibridge.CopilotConfig{
|
||||
Name: row.Name,
|
||||
BaseURL: row.BaseUrl,
|
||||
Name: spec.Name,
|
||||
BaseURL: spec.BaseURL,
|
||||
APIDumpDir: dumpDir,
|
||||
CircuitBreaker: cbCfg,
|
||||
}), nil
|
||||
|
||||
default:
|
||||
return nil, xerrors.Errorf("unsupported provider type: %q", row.Type)
|
||||
return nil, xerrors.Errorf("unsupported provider type: %q", spec.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// disabledProviderFromRow builds a Provider stub for a disabled row.
|
||||
// Using provider.DisabledStub rather than a concrete provider avoids
|
||||
// duplicating the row.Type switch and ensures that a new AIProviderType
|
||||
// value is automatically handled without requiring a matching case here.
|
||||
func disabledProviderFromRow(row database.AIProvider) (aibridge.Provider, error) {
|
||||
return aibridge.NewDisabledProviderStub(row.Name, string(row.Type)), nil
|
||||
}
|
||||
|
||||
// buildAIProviderKeyPool builds a [keypool.Pool]. Callers must check
|
||||
// len(keys) > 0 first; keypool.New rejects empty input.
|
||||
func buildAIProviderKeyPool(providerName string, keys []database.AIProviderKey, metrics *aibridge.Metrics) (*keypool.Pool, error) {
|
||||
raw := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
raw = append(raw, k.APIKey)
|
||||
}
|
||||
return keypool.New(providerName, raw, quartz.NewReal(), metrics)
|
||||
func buildAIProviderKeyPool(providerName string, keys []string, metrics *aibridge.Metrics) (*keypool.Pool, error) {
|
||||
return keypool.New(providerName, keys, quartz.NewReal(), metrics)
|
||||
}
|
||||
|
||||
// bedrockConfigFromRow returns nil when the settings have no Bedrock
|
||||
// discriminator or when the Bedrock fields are not actually configured.
|
||||
// The provider row's BaseUrl is the generic upstream endpoint and is
|
||||
// always non-empty, so it cannot serve as a Bedrock detection signal;
|
||||
// gate on the settings blob alone via [codersdk.AIProviderBedrockSettings.IsConfigured].
|
||||
func bedrockConfigFromRow(row database.AIProvider, settings codersdk.AIProviderSettings) *aibridge.AWSBedrockConfig {
|
||||
if settings.Bedrock == nil {
|
||||
// bedrockConfig returns nil when the settings are absent or when the
|
||||
// Bedrock fields are not actually configured. The provider's BaseURL is
|
||||
// the generic upstream endpoint and is always non-empty, so it cannot
|
||||
// serve as a Bedrock detection signal; gate on the settings alone via
|
||||
// [codersdk.AIProviderBedrockSettings.IsConfigured].
|
||||
func bedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) *aibridge.AWSBedrockConfig {
|
||||
if bedrock == nil {
|
||||
return nil
|
||||
}
|
||||
bedrockSettings := *settings.Bedrock
|
||||
bedrockSettings := *bedrock
|
||||
if !bedrockSettings.IsConfigured() {
|
||||
return nil
|
||||
}
|
||||
accessKey := ptr.NilToEmpty(bedrockSettings.AccessKey)
|
||||
accessKeySecret := ptr.NilToEmpty(bedrockSettings.AccessKeySecret)
|
||||
return &aibridge.AWSBedrockConfig{
|
||||
BaseURL: row.BaseUrl,
|
||||
BaseURL: baseURL,
|
||||
Region: bedrockSettings.Region,
|
||||
AccessKey: accessKey,
|
||||
AccessKeySecret: accessKeySecret,
|
||||
|
||||
Reference in New Issue
Block a user