| 325 | } |
| 326 | |
| 327 | func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) { |
| 328 | //nolint:gocritic // AIBridged has specific authz rules. |
| 329 | ctx = dbauthz.AsAIBridged(ctx) |
| 330 | |
| 331 | intcID, err := uuid.Parse(in.GetInterceptionId()) |
| 332 | if err != nil { |
| 333 | return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) |
| 334 | } |
| 335 | |
| 336 | metadata := metadataToMap(in.GetMetadata()) |
| 337 | |
| 338 | if s.structuredLogging { |
| 339 | s.logger.Info(ctx, InterceptionLogMarker, |
| 340 | slog.F("record_type", "tool_usage"), |
| 341 | slog.F("interception_id", intcID.String()), |
| 342 | slog.F("msg_id", in.GetMsgId()), |
| 343 | slog.F("tool_call_id", in.GetToolCallId()), |
| 344 | slog.F("tool", in.GetTool()), |
| 345 | slog.F("input", in.GetInput()), |
| 346 | slog.F("server_url", in.GetServerUrl()), |
| 347 | slog.F("injected", in.GetInjected()), |
| 348 | slog.F("invocation_error", in.GetInvocationError()), |
| 349 | slog.F("created_at", in.GetCreatedAt().AsTime()), |
| 350 | slog.F("metadata", metadata), |
| 351 | ) |
| 352 | } |
| 353 | |
| 354 | out, err := json.Marshal(metadata) |
| 355 | if err != nil { |
| 356 | s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err)) |
| 357 | } |
| 358 | |
| 359 | _, err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{ |
| 360 | ID: uuid.New(), |
| 361 | InterceptionID: intcID, |
| 362 | ProviderResponseID: in.GetMsgId(), |
| 363 | ProviderToolCallID: sql.NullString{String: in.GetToolCallId(), Valid: in.GetToolCallId() != ""}, |
| 364 | ServerUrl: sql.NullString{String: in.GetServerUrl(), Valid: in.ServerUrl != nil}, |
| 365 | Tool: in.GetTool(), |
| 366 | Input: in.GetInput(), |
| 367 | Injected: in.GetInjected(), |
| 368 | InvocationError: sql.NullString{String: in.GetInvocationError(), Valid: in.InvocationError != nil}, |
| 369 | Metadata: out, |
| 370 | CreatedAt: in.GetCreatedAt().AsTime(), |
| 371 | }) |
| 372 | if err != nil { |
| 373 | return nil, xerrors.Errorf("insert tool usage: %w", err) |
| 374 | } |
| 375 | |
| 376 | return &proto.RecordToolUsageResponse{}, nil |
| 377 | } |
| 378 | |
| 379 | func (s *Server) RecordModelThought(ctx context.Context, in *proto.RecordModelThoughtRequest) (*proto.RecordModelThoughtResponse, error) { |
| 380 | //nolint:gocritic // AIBridged has specific authz rules. |