org.wildfly.extension.ai.chat.ChatModelConnectionCheckerOperationHandler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of wildfly-ai-subsystem Show documentation
There is a newer version: 0.1.0
/*
 * Copyright The WildFly Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package org.wildfly.extension.ai.chat;

import static org.jboss.dmr.ModelType.STRING;
import static org.wildfly.extension.ai.AIAttributeDefinitions.USER_MESSAGE;

import dev.langchain4j.model.chat.ChatLanguageModel;
import org.jboss.as.controller.OperationContext;
import org.jboss.as.controller.OperationFailedException;
import org.jboss.as.controller.OperationStepHandler;
import org.jboss.as.controller.SimpleOperationDefinitionBuilder;
import org.jboss.as.controller.access.management.SensitiveTargetAccessConstraintDefinition;
import org.jboss.as.controller.registry.ManagementResourceRegistration;
import org.jboss.dmr.ModelNode;
import org.wildfly.extension.ai.AILogger;
import org.wildfly.service.capture.FunctionExecutor;
import org.wildfly.service.capture.ValueExecutorRegistry;
import org.wildfly.subsystem.resource.ResourceDescriptor;

/**
 *
 * @author Emmanuel Hugonnet (c) 2024 Red Hat, Inc.
 */
public class ChatModelConnectionCheckerOperationHandler implements OperationStepHandler {
    private static final String OPERATION_NAME = "chat";

    public static void register(final ManagementResourceRegistration registration, final ResourceDescriptor descriptor, ValueExecutorRegistry registry) {
        registration.registerOperationHandler(SimpleOperationDefinitionBuilder.of(OPERATION_NAME, descriptor.getResourceDescriptionResolver())
                .setParameters(USER_MESSAGE)
                .setReplyType(STRING)
                .setRuntimeOnly()
                .setReadOnly()
                .addAccessConstraint(SensitiveTargetAccessConstraintDefinition.READ_WHOLE_CONFIG)
                .build(), new ChatModelConnectionCheckerOperationHandler(registry));
    }

    private final ValueExecutorRegistry registry;

    ChatModelConnectionCheckerOperationHandler(ValueExecutorRegistry registry) {
        this.registry = registry;
    }

    @Override
    public void execute(OperationContext context, ModelNode operation) throws OperationFailedException {
        String userMessage = USER_MESSAGE.resolveModelAttribute(context, operation).asString();
        FunctionExecutor executor = registry.getExecutor(context.getCurrentAddressValue());
        if(executor == null) {
            throw AILogger.ROOT_LOGGER.chatLanguageModelServiceUnavailable(context.getCurrentAddressValue());
        }
        ModelNode answer = executor.execute((ChatLanguageModel chatLanguageModel) -> {
            AILogger.ROOT_LOGGER.debug("About to execute a chat call to the LLM with the following user message: " + userMessage);
            String response = chatLanguageModel.generate(userMessage);
            AILogger.ROOT_LOGGER.debug("This is the answer I got: " + response);
            return new ModelNode(response);
        });
        context.getResult().set(answer);
    }

}