coder/enterprise/cli/aibridged.go

//go:build !slim

package cli

import (
	"context"

	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/xerrors"

	"github.com/coder/aibridge"
	"github.com/coder/aibridge/config"
	agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
	"github.com/coder/coder/v2/coderd/tracing"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/coder/v2/enterprise/aibridged"
	"github.com/coder/coder/v2/enterprise/coderd"
)

func newAIBridgeDaemon(coderAPI *coderd.API) (*aibridged.Server, error) {
	ctx := context.Background()
	coderAPI.Logger.Debug(ctx, "starting in-memory aibridge daemon")

	logger := coderAPI.Logger.Named("aibridged")
	cfg := coderAPI.DeploymentValues.AI.BridgeConfig

	// Build circuit breaker config if enabled.
	var cbConfig *config.CircuitBreaker
	if cfg.CircuitBreakerEnabled.Value() {
		cbConfig = &config.CircuitBreaker{
			FailureThreshold: uint32(cfg.CircuitBreakerFailureThreshold.Value()), //nolint:gosec // Validated by serpent.Validate in deployment options.
			Interval:         cfg.CircuitBreakerInterval.Value(),
			Timeout:          cfg.CircuitBreakerTimeout.Value(),
			MaxRequests:      uint32(cfg.CircuitBreakerMaxRequests.Value()), //nolint:gosec // Validated by serpent.Validate in deployment options.
		}
	}

	// Setup supported providers with circuit breaker config.
	providers := []aibridge.Provider{
		aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
			Name:             aibridge.ProviderOpenAI,
			BaseURL:          cfg.OpenAI.BaseURL.String(),
			Key:              cfg.OpenAI.Key.String(),
			CircuitBreaker:   cbConfig,
			SendActorHeaders: cfg.SendActorHeaders.Value(),
		}),
		aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
			Name:             aibridge.ProviderAnthropic,
			BaseURL:          cfg.Anthropic.BaseURL.String(),
			Key:              cfg.Anthropic.Key.String(),
			CircuitBreaker:   cbConfig,
			SendActorHeaders: cfg.SendActorHeaders.Value(),
		}, getBedrockConfig(cfg.Bedrock)),
		aibridge.NewCopilotProvider(aibridge.CopilotConfig{
			Name:           aibridge.ProviderCopilot,
			CircuitBreaker: cbConfig,
		}),
		aibridge.NewCopilotProvider(aibridge.CopilotConfig{
			Name:           agplaibridge.ProviderCopilotBusiness,
			BaseURL:        "https://" + agplaibridge.HostCopilotBusiness,
			CircuitBreaker: cbConfig,
		}),
		aibridge.NewCopilotProvider(aibridge.CopilotConfig{
			Name:           agplaibridge.ProviderCopilotEnterprise,
			BaseURL:        "https://" + agplaibridge.HostCopilotEnterprise,
			CircuitBreaker: cbConfig,
		}),
		aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
			Name:             agplaibridge.ProviderChatGPT,
			BaseURL:          agplaibridge.BaseURLChatGPT,
			CircuitBreaker:   cbConfig,
			SendActorHeaders: cfg.SendActorHeaders.Value(),
		}),
	}

	reg := prometheus.WrapRegistererWithPrefix("coder_aibridged_", coderAPI.PrometheusRegistry)
	metrics := aibridge.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.
	if err != nil {
		return nil, xerrors.Errorf("create request pool: %w", err)
	}

	// Create daemon.
	srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
		return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
	}, logger, tracer)
	if err != nil {
		return nil, xerrors.Errorf("start in-memory aibridge daemon: %w", err)
	}
	return srv, nil
}

func getBedrockConfig(cfg codersdk.AIBridgeBedrockConfig) *aibridge.AWSBedrockConfig {
	if cfg.Region.String() == "" && cfg.BaseURL.String() == "" && cfg.AccessKey.String() == "" && cfg.AccessKeySecret.String() == "" {
		return nil
	}

	return &aibridge.AWSBedrockConfig{
		BaseURL:         cfg.BaseURL.String(),
		Region:          cfg.Region.String(),
		AccessKey:       cfg.AccessKey.String(),
		AccessKeySecret: cfg.AccessKeySecret.String(),
		Model:           cfg.Model.String(),
		SmallFastModel:  cfg.SmallFastModel.String(),
	}
}