diff --git a/aibridge/recorder/recorder.go b/aibridge/recorder/recorder.go index 29c1ccdafd876..98aece3c98ce4 100644 --- a/aibridge/recorder/recorder.go +++ b/aibridge/recorder/recorder.go @@ -18,7 +18,7 @@ var ( _ Recorder = &AsyncRecorder{} ) -// WrappedRecorder is a convenience struct which implements RecorderClient and resolves a client before calling each method. +// WrappedRecorder is a convenience struct which implements Recorder and resolves a client before calling each method. // It also sets the start/creation time of each record. type WrappedRecorder struct { logger slog.Logger diff --git a/coderd/aibridged/translator.go b/coderd/aibridged/drpc_recorder.go similarity index 84% rename from coderd/aibridged/translator.go rename to coderd/aibridged/drpc_recorder.go index 536c00bcabc9e..9ad48997311c5 100644 --- a/coderd/aibridged/translator.go +++ b/coderd/aibridged/drpc_recorder.go @@ -14,15 +14,15 @@ import ( "github.com/coder/coder/v2/coderd/aibridged/proto" ) -var _ aibridge.Recorder = &recorderTranslation{} +var _ aibridge.Recorder = &DRPCRecorder{} -// recorderTranslation satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver. -type recorderTranslation struct { +// DRPCRecorder satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver. +type DRPCRecorder struct { apiKeyID string client proto.DRPCRecorderClient } -func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error { +func (t *DRPCRecorder) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error { _, err := t.client.RecordInterception(ctx, &proto.RecordInterceptionRequest{ Id: req.ID, ApiKeyId: t.apiKeyID, @@ -44,7 +44,7 @@ func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibri return err } -func (t *recorderTranslation) RecordInterceptionEnded(ctx context.Context, req *aibridge.InterceptionRecordEnded) error { +func (t *DRPCRecorder) RecordInterceptionEnded(ctx context.Context, req *aibridge.InterceptionRecordEnded) error { endedReq := &proto.RecordInterceptionEndedRequest{ Id: req.ID, EndedAt: timestamppb.New(req.EndedAt), @@ -61,7 +61,7 @@ func (t *recorderTranslation) RecordInterceptionEnded(ctx context.Context, req * return err } -func (t *recorderTranslation) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error { +func (t *DRPCRecorder) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error { _, err := t.client.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{ InterceptionId: req.InterceptionID, MsgId: req.MsgID, @@ -72,7 +72,7 @@ func (t *recorderTranslation) RecordPromptUsage(ctx context.Context, req *aibrid return err } -func (t *recorderTranslation) RecordTokenUsage(ctx context.Context, req *aibridge.TokenUsageRecord) error { +func (t *DRPCRecorder) RecordTokenUsage(ctx context.Context, req *aibridge.TokenUsageRecord) error { merged := req.Metadata if merged == nil { merged = aibridge.Metadata{} @@ -96,7 +96,7 @@ func (t *recorderTranslation) RecordTokenUsage(ctx context.Context, req *aibridg return err } -func (t *recorderTranslation) RecordToolUsage(ctx context.Context, req *aibridge.ToolUsageRecord) error { +func (t *DRPCRecorder) RecordToolUsage(ctx context.Context, req *aibridge.ToolUsageRecord) error { serialized, err := json.Marshal(req.Args) if err != nil { return xerrors.Errorf("serialize tool %q args: %w", req.Tool, err) @@ -123,7 +123,7 @@ func (t *recorderTranslation) RecordToolUsage(ctx context.Context, req *aibridge return err } -func (t *recorderTranslation) RecordModelThought(ctx context.Context, req *aibridge.ModelThoughtRecord) error { +func (t *DRPCRecorder) RecordModelThought(ctx context.Context, req *aibridge.ModelThoughtRecord) error { _, err := t.client.RecordModelThought(ctx, &proto.RecordModelThoughtRequest{ InterceptionId: req.InterceptionID, Content: req.Content, diff --git a/coderd/aibridged/pool.go b/coderd/aibridged/pool.go index 3eb3630a5de3e..5649dedf7bcd9 100644 --- a/coderd/aibridged/pool.go +++ b/coderd/aibridged/pool.go @@ -228,16 +228,20 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl span.AddEvent("cache_miss") providerVersion := p.providerVersion.Load() - recorder := aibridge.NewRecorder(p.logger.Named("recorder"), p.tracer, func(clientCtx context.Context) (aibridge.Recorder, error) { - // The recorder outlives this Acquire call, so the client is acquired - // against the context of the record call being served. - client, err := clientFn(clientCtx) - if err != nil { - return nil, xerrors.Errorf("acquire client: %w", err) - } + recorder := aibridge.NewRecorder( + p.logger.Named("recorder"), + p.tracer, + func(clientCtx context.Context) (aibridge.Recorder, error) { + // The recorder outlives this Acquire call, so the client is acquired + // against the context of the record call being served. + client, err := clientFn(clientCtx) + if err != nil { + return nil, xerrors.Errorf("acquire client: %w", err) + } - return &recorderTranslation{apiKeyID: req.APIKeyID, client: client}, nil - }) + return &DRPCRecorder{apiKeyID: req.APIKeyID, client: client}, nil + }, + ) // Slow path. // Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.