mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
Paweł Banaszewski
ffc7dadfe2
Adds typed AI Gateway actor identity, including the API key ID used by recording. Recorders now read the authenticated API key ID from each call context instead of capturing request identity at construction time. This allows one recorder to be reused safely across sequential and concurrent requests. Includes context helper coverage and recorder regression tests. Generated by Coder Agents on behalf of @pawbana.
105 lines
4.0 KiB
Go
105 lines
4.0 KiB
Go
package aibridge
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"github.com/coder/coder/v2/aibridge/config"
|
|
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
|
"github.com/coder/coder/v2/aibridge/metrics"
|
|
"github.com/coder/coder/v2/aibridge/provider"
|
|
"github.com/coder/coder/v2/aibridge/recorder"
|
|
)
|
|
|
|
// Const + Type + function aliases for backwards compatibility.
|
|
const (
|
|
ProviderAnthropic = config.ProviderAnthropic
|
|
ProviderBedrock = config.ProviderBedrock
|
|
ProviderOpenAI = config.ProviderOpenAI
|
|
ProviderCopilot = config.ProviderCopilot
|
|
)
|
|
|
|
type (
|
|
Metrics = metrics.Metrics
|
|
|
|
Provider = provider.Provider
|
|
|
|
// Actor is the authenticated identity of an AI Gateway request.
|
|
Actor = aibcontext.Actor
|
|
|
|
InterceptionRecord = recorder.InterceptionRecord
|
|
InterceptionRecordEnded = recorder.InterceptionRecordEnded
|
|
TokenUsageRecord = recorder.TokenUsageRecord
|
|
PromptUsageRecord = recorder.PromptUsageRecord
|
|
ToolUsageRecord = recorder.ToolUsageRecord
|
|
ModelThoughtRecord = recorder.ModelThoughtRecord
|
|
Recorder = recorder.Recorder
|
|
Metadata = recorder.Metadata
|
|
ErrorType = recorder.ErrorType
|
|
|
|
AnthropicConfig = config.Anthropic
|
|
AWSBedrockConfig = config.AWSBedrock
|
|
AWSClaudePlatformConfig = config.AWSClaudePlatform
|
|
OpenAIConfig = config.OpenAI
|
|
CopilotConfig = config.Copilot
|
|
)
|
|
|
|
// AsActor attaches the authenticated identity to an AI Gateway request context.
|
|
func AsActor(ctx context.Context, actor Actor) context.Context {
|
|
return aibcontext.AsActor(ctx, actor)
|
|
}
|
|
|
|
// NewAnthropicProvider constructs the Anthropic provider. At most one of
|
|
// bedrockCfg and claudePlatformCfg may be non-nil; both nil means direct
|
|
// Anthropic API access.
|
|
func NewAnthropicProvider(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock, claudePlatformCfg *config.AWSClaudePlatform) (provider.Provider, error) {
|
|
return provider.NewAnthropic(ctx, cfg, bedrockCfg, claudePlatformCfg)
|
|
}
|
|
|
|
func NewBedrockProvider(ctx context.Context, cfg config.Anthropic, bedrockCfg config.AWSBedrock) (provider.Provider, error) {
|
|
return provider.NewBedrock(ctx, cfg, bedrockCfg)
|
|
}
|
|
|
|
func NewOpenAIProvider(cfg config.OpenAI) provider.Provider {
|
|
return provider.NewOpenAI(cfg)
|
|
}
|
|
|
|
func NewCopilotProvider(cfg config.Copilot) provider.Provider {
|
|
return provider.NewCopilot(cfg)
|
|
}
|
|
|
|
// NewDisabledProviderStub returns a Provider that reports Enabled() ==
|
|
// false and has no-op implementations for all other methods. Use this
|
|
// instead of constructing a concrete provider for disabled rows so that
|
|
// adding a new provider type does not require updating a switch here.
|
|
func NewDisabledProviderStub(name, providerType string) provider.Provider {
|
|
return provider.NewDisabledStub(name, providerType)
|
|
}
|
|
|
|
func NewMetrics(reg prometheus.Registerer) *metrics.Metrics {
|
|
return metrics.NewMetrics(reg)
|
|
}
|
|
|
|
// NewRecorder creates a [Recorder] which logs each record and refuses
|
|
// malformed ones before handing it to base. Interception start records require
|
|
// an actor with an API-key ID on their context.
|
|
//
|
|
// middleware is inserted below the logging and validating middleware, so that
|
|
// every record is logged and checked before any of it runs, and above the
|
|
// tracing middleware, which must stay immediately above the recorder its spans
|
|
// measure. Policy that drops records, such as [recorder.WithoutRecords],
|
|
// belongs here: it keeps NewRecorder to its own concerns and leaves the choice
|
|
// to the caller.
|
|
func NewRecorder(logger slog.Logger, tracer trace.Tracer, structured bool, base Recorder, middleware ...recorder.Middleware) Recorder {
|
|
chain := make([]recorder.Middleware, 0, len(middleware)+3)
|
|
chain = append(chain, recorder.WithLogging(logger, structured))
|
|
chain = append(chain, recorder.WithValidation(logger))
|
|
chain = append(chain, middleware...)
|
|
chain = append(chain, recorder.WithTracing(tracer))
|
|
|
|
return recorder.ChainMiddleware(chain...)(base)
|
|
}
|