| 235 | } |
| 236 | |
| 237 | func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) { |
| 238 | //nolint:gocritic // AIBridged has specific authz rules. |
| 239 | ctx = dbauthz.AsAIBridged(ctx) |
| 240 | |
| 241 | intcID, err := uuid.Parse(in.GetInterceptionId()) |
| 242 | if err != nil { |
| 243 | return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) |
| 244 | } |
| 245 | |
| 246 | metadata := metadataToMap(in.GetMetadata()) |
| 247 | |
| 248 | if s.structuredLogging { |
| 249 | s.logger.Info(ctx, InterceptionLogMarker, |
| 250 | slog.F("record_type", "token_usage"), |
| 251 | slog.F("interception_id", intcID.String()), |
| 252 | slog.F("msg_id", in.GetMsgId()), |
| 253 | slog.F("input_tokens", in.GetInputTokens()), |
| 254 | slog.F("output_tokens", in.GetOutputTokens()), |
| 255 | slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()), |
| 256 | slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()), |
| 257 | slog.F("created_at", in.GetCreatedAt().AsTime()), |
| 258 | slog.F("metadata", metadata), |
| 259 | ) |
| 260 | } |
| 261 | |
| 262 | out, err := json.Marshal(metadata) |
| 263 | if err != nil { |
| 264 | s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err)) |
| 265 | } |
| 266 | |
| 267 | _, err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{ |
| 268 | ID: uuid.New(), |
| 269 | InterceptionID: intcID, |
| 270 | ProviderResponseID: in.GetMsgId(), |
| 271 | InputTokens: in.GetInputTokens(), |
| 272 | OutputTokens: in.GetOutputTokens(), |
| 273 | CacheReadInputTokens: in.GetCacheReadInputTokens(), |
| 274 | CacheWriteInputTokens: in.GetCacheWriteInputTokens(), |
| 275 | Metadata: out, |
| 276 | CreatedAt: in.GetCreatedAt().AsTime(), |
| 277 | }) |
| 278 | if err != nil { |
| 279 | return nil, xerrors.Errorf("insert token usage: %w", err) |
| 280 | } |
| 281 | |
| 282 | return &proto.RecordTokenUsageResponse{}, nil |
| 283 | } |
| 284 | |
| 285 | func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) { |
| 286 | //nolint:gocritic // AIBridged has specific authz rules. |