@@ -15,7 +15,6 @@ import (
1515"golang.org/x/xerrors"
1616
1717"github.com/coder/coder/v2/aibridged/proto"
18- "github.com/coder/coder/v2/coderd/database"
1918"github.com/coder/coder/v2/codersdk"
2019)
2120
@@ -49,16 +48,17 @@ type Server struct {
4948// shuttingDownCh will receive when we start graceful shutdown
5049shuttingDownCh chan struct {}
5150
52- bridge * Bridge
51+ bridge * Bridge
5352}
5453
55- func New (store database.Store ,rpcDialer Dialer ,httpAddr string ,logger slog.Logger ) (* Server ,error ) {
54+ var _ proto.DRPCAIBridgeDaemonServer = & Server {}
55+
56+ func New (rpcDialer Dialer ,httpAddr string ,logger slog.Logger ) (* Server ,error ) {
5657if rpcDialer == nil {
5758return nil ,xerrors .Errorf ("nil rpcDialer given" )
5859}
5960
6061ctx ,cancel := context .WithCancel (context .Background ())
61- bridge := NewBridge (httpAddr ,store )
6262daemon := & Server {
6363logger :logger ,
6464clientDialer :rpcDialer ,
@@ -68,9 +68,11 @@ func New(store database.Store, rpcDialer Dialer, httpAddr string, logger slog.Lo
6868closedCh :make (chan struct {}),
6969shuttingDownCh :make (chan struct {}),
7070initConnectionCh :make (chan struct {}),
71-
72- bridge :bridge ,
7371}
72+
73+ bridge := NewBridge (httpAddr ,daemon .client )
74+ daemon .bridge = bridge
75+
7476go daemon .connect ()
7577go func () {
7678err := bridge .Serve ()
@@ -164,6 +166,26 @@ func (s *Server) AuditPrompt(ctx context.Context, in *proto.AuditPromptRequest)
164166return out ,nil
165167}
166168
169+ func (s * Server )TrackTokenUsage (ctx context.Context ,in * proto.TrackTokenUsageRequest ) (* proto.TrackTokenUsageResponse ,error ) {
170+ out ,err := clientDoWithRetries (ctx ,s .client ,func (ctx context.Context ,client proto.DRPCAIBridgeDaemonClient ) (* proto.TrackTokenUsageResponse ,error ) {
171+ return client .TrackTokenUsage (ctx ,in )
172+ })
173+ if err != nil {
174+ return nil ,err
175+ }
176+ return out ,nil
177+ }
178+
179+ func (s * Server )TrackUserPrompts (ctx context.Context ,in * proto.TrackUserPromptsRequest ) (* proto.TrackUserPromptsResponse ,error ) {
180+ out ,err := clientDoWithRetries (ctx ,s .client ,func (ctx context.Context ,client proto.DRPCAIBridgeDaemonClient ) (* proto.TrackUserPromptsResponse ,error ) {
181+ return client .TrackUserPrompts (ctx ,in )
182+ })
183+ if err != nil {
184+ return nil ,err
185+ }
186+ return out ,nil
187+ }
188+
167189//func (s *Server) ChatCompletions(payload *proto.JSONPayload, stream proto.DRPCOpenAIService_ChatCompletionsStream) error {
168190//// TODO: call OpenAI API.
169191//