Spring AI Adding custom CardinalityKeyValue did not work

I want to add some custom content to the observable log, such as user ID, request ID, etc., but it doesn't work. Is it because I used it incorrectly, How should it be implemented?

In the following example, the ChatModelObservationHandler always fails to retrieve the user ID.

public class SpringAIController {

    @Resource
    private OpenAiChatModel openAiChatModel;
    @Resource
    private ObservationRegistry observationRegistry;

    @PostMapping("/api/test/openai_example")
    public ResponseEntity<?> testOpenAI(@RequestParam(required = false, defaultValue = "你好") String content) {
        String userId = "test-user";

        Objects.requireNonNull(observationRegistry.getCurrentObservation()).lowCardinalityKeyValue("userId", userId);

        String response = openAiChatModel.call(content);
        return ResponseEntity.ok(response);
    }

}



@Component
public class ChatModelObservationHandler implements ObservationHandler<ChatModelObservationContext> {

    @Override
    public void onStart(ChatModelObservationContext context) {
        Prompt request = context.getRequest();
        String userId = context.getLowCardinalityKeyValues().stream()
                .filter(kv -> "userId".equals(kv.getKey()))
                .map(KeyValue::getValue)
                .findFirst()
                .orElse("unknown");
        log.info("ChatModelObservationHandler onStart: userId={}, request={}", userId, request);
    }

    @Override
    public void onStop(ChatModelObservationContext context) {
        Usage usage = usage(context);
        List<String> completions = completion(context);
        Prompt request = context.getRequest();
        String userId = context.getLowCardinalityKeyValues().stream()
                .filter(kv -> "userId".equals(kv.getKey()))
                .map(KeyValue::getValue)
                .findFirst()
                .orElse("unknown");
        log.info("ChatModelObservationHandler onStop: userId={}, request={}, completions={}, usage={}", userId,
                request, completions, usage);
    }

    @Override
    public boolean supportsContext(Observation.Context context) {
        return context instanceof ChatModelObservationContext;
    }

    private List<String> completion(ChatModelObservationContext context) {
        if (context.getResponse() == null || context.getResponse().getResults() == null
                || CollectionUtils.isEmpty(context.getResponse().getResults())) {
            return List.of();
        }

        if (!StringUtils.hasText(context.getResponse().getResult().getOutput().getText())) {
            return List.of();
        }

        return context.getResponse()
                .getResults()
                .stream()
                .filter(generation -> generation.getOutput() != null
                        && StringUtils.hasText(generation.getOutput().getText()))
                .map(generation -> generation.getOutput().getText())
                .toList();
    }

    private Usage usage(ChatModelObservationContext context) {
        if (context.getResponse() == null || context.getResponse().getMetadata() == null
                || context.getResponse().getMetadata().getUsage() == null) {
            return null;
        }
        return context.getResponse().getMetadata().getUsage();
    }

}

Comment From: WOONBE

If this issue is not yet resolved on your end, the correct way to implement this is to have the ChatModelObservationHandler explicitly access its parent observation to retrieve the userId. The context provides a getParentObservation() method for this exact purpose.

Here's an example code to retrieve userId from ChatModelObservationContext.

public String findUserIdFromParent(ChatModelObservationContext context) {

            Observation parentObservation = (Observation) context.getParentObservation();
            if (parentObservation == null) {
                return "unknown";
            }

            return parentObservation.getContext().getLowCardinalityKeyValues().stream()
                    .filter(kv -> "userId".equals(kv.getKey()))
                    .map(KeyValue::getValue)
                    .findFirst()
                    .orElse("unknown");
        }

I hope this explanation helps you to resolve the problem.