AI Module Extensibility Guide

The AI module provides a number of extensibility points, and custom implementations can be used for:

• Model Providers
• Vector Stores (experimental)
• Tools Calling

Model Providers

Interface com.bloomreach.xm.ai.service.api.client.AIModelFactory can be implemented to wire a different org.springframework.ai.chat.model.ChatModel and org.springframework.ai.embedding.EmbeddingMode:

public interface AIModelFactory<C extends ChatModel, E extends EmbeddingModel> {
    String getModelProviderName();
    boolean isEnabled();
    List<ModelFeature> getModelFeatures();
    C createChatModel();
    E createEmbeddingModel();
}

Model provider factories in the AI backend are discovered and injected dynamically (Spring Dependency Injection), as long as they are annotated with Spring's @Component or @Service and are visible in the classpath.

In order to have access to (both JCR and Properties) configuration, it's better to start from the abstract com.bloomreach.xm.ai.service.impl.client.chatmodel.JcrAIModelFactory. Here are parts of the Ollama provider:

@Service
public class OllamaModelFactory extends JcrAIModelFactory<OllamaChatModel, OllamaEmbeddingModel> {

public OllamaModelFactory(final JcrFirstConfiguration jcrFirstConfiguration) {
    super(jcrFirstConfiguration);

    if (isEnabled()) {
        List.of(SPRING_AI_OLLAMA_API_URL_PROP_NAME,  ...
                SPRING_AI_OLLAMA_PULL_STRATEGY_PROP_NAME)
                .forEach(this::ensureRequiredPropertyConfigured);
    }

    if (isEnabled()) {
        log.info("ChatModelFactory [{}] active", getModelProviderName());
    }
}

@Override
public String getModelProviderName() {
    return "Ollama";
}

@Override
public OllamaChatModel createChatModel() {
    if (!isEnabled()) {
        throw new IllegalStateException("Ollama Model is not enabled");
    }

    return OllamaChatModel.builder()
            .ollamaApi(OllamaApi.builder()
                    .baseUrl(getProperty(SPRING_AI_OLLAMA_API_URL_PROP_NAME)).build())
            .defaultOptions(OllamaChatOptions.builder()
                    .model(getProperty(SPRING_AI_OLLAMA_CHAT_OPTIONS_MODEL_PROP_NAME))
                    .build())
             ...
            .build();
}

@Override
public OllamaEmbeddingModel createEmbeddingModel() {
    if (!isEnabled()) {
        throw new IllegalStateException("Ollama Model is not enabled");
    }

    final OllamaEmbeddingOptions.Builder optionsBuilder = OllamaEmbeddingOptions.builder();
    if (isConfigured(SPRING_AI_OLLAMA_EMBEDDING_OPTIONS_MODEL_PROP_NAME)) {
        optionsBuilder.model(getProperty(SPRING_AI_OLLAMA_EMBEDDING_OPTIONS_MODEL_PROP_NAME));
    }

    return new OllamaEmbeddingModel(
            OllamaApi.builder().baseUrl(getProperty(SPRING_AI_OLLAMA_API_URL_PROP_NAME)).build(),
            optionsBuilder.build(),
            ObservationRegistry.NOOP,
            ...
}

Vector Stores

Interface com.bloomreach.xm.ai.service.api.client.VectorStoreFactory can be used to implement custom vector store wiring, so that a different org.springframework.ai.vectorstore.VectorStore can be used:

public interface VectorStoreFactory<V extends VectorStore> {
    String getVectorStoreName();
    boolean isEnabled();
    V createVectorStore();
}

Vector store factories in the AI backend are discovered and injected dynamically (Spring Dependency Injection), as long as they are annotated with Spring's @Component or @Service and are visible in the classpath.

In order to have access to (both JCR and Properties) configuration, it's better to start from the abstract com.bloomreach.xm.ai.service.impl.vector.JcrVectorStoreFactory. Here are parts of the Redis factory:

@Service
public class RedisVectorStoreFactory extends JcrVectorStoreFactory<RedisVectorStore> {

    public RedisVectorStoreFactory(final JcrFirstConfiguration jcrFirstConfiguration,
                                   final ObjectProvider<EmbeddingModel> embeddingModel) {
        super(jcrFirstConfiguration);

        this.embeddingModel = embeddingModel.getIfAvailable();
        if (this.embeddingModel == null) {
            log.warn("EmbeddingModel should not be null, disabling store");
            disabled = true;
        }

        if (isEnabled()) {
            List.of(REDIS_HOST_PROP_NAME, REDIS_PORT_PROP_NAME,
                     ...
                    .forEach(this::ensureRequiredPropertyConfigured);
        }

        if (isEnabled()) {
            log.info("VectorStoreFactory [{}] active", getVectorStoreName());
        }
    }

    @Override
    public String getVectorStoreName() {
        return "Redis";
    }

    @Override
    public RedisVectorStore createVectorStore() {
        if (!isEnabled()) {
            throw new IllegalStateException("Redis Vector Store is not enabled");
        }

        return RedisVectorStore.builder(...)
                .indexName(getProperty(REDIS_INDEX_PROP_NAME))
                .prefix(getProperty(REDIS_PREFIX_PROP_NAME))
                ...
                .build();
    }
}

 

Tool Packages

Interface com.bloomreach.xm.ai.service.api.client.ClientToolPackage allows registration of custom tools that the AI is made aware of and can call. The interface allows registering a tool package, thus a container for more than a single tool.

public interface ClientToolPackage {
    boolean isEnabled();
    String getAdvice();
    Object getTools();
} 

Prompt instructions

Along the getTools() method, the interface defines method getAdvice(), which is used for appending natural language instructions, regarding the tools in your package, to the system prompt. The instructions are meant to help explain what the tool does. These instructions can alternatively be provided in the tools descriptions (see the @Tool annotation in example code), thus an empty string is also a valid return value for method getAdvice().

Registration

Tools in the AI backend are discovered and injected dynamically (Spring Dependency Injection), as long as they are annotated with Spring's @Component or @Service and are visible in the classpath. This immediately enables building and registering custom tools for end projects.

Capabilities

The AI module's Spring beans are available to the tools. For example, to obtain the Vector Store, all you need is a constructor argument: ObjectProvider<VectorStore> vectorStore.

Tools are executed during a user's chat. The tools are executed among this sequence of events: User chats -> AI receives request -> AI calls tool -> tool executes and responds to the AI -> AI responds to user.

The order in which tools are executed can be important. The order is specified with the Order annotation, eg @Order(10).

Tools also receive "context", this is a Map that contains attributes set by the tools caller (the AI backend). Accessing items from the context is easy via the helper methods of com.bloomreach.xm.ai.service.impl.client.advisors.AdvisorContextUtils:

static <T> Optional<T> getFromContext(final Map<String, Object> context, final String key, final Class<T> type)

static Optional<String> getStringFromContext(final Map<String, Object> context, final String key)

The available attributes are:

• userId: The username of the user that is chatting with the AI (String)
• activeContextItem: An object of type com.bloomreach.xm.ai.repository.services.content.model.Document that represents the current open document the chatting user is working on

Tools should be as fast as possible and they should be blocking, the AI (and the user) is waiting for their response.

Example Tool Package

Below is an example of a tool package that exposes two tools:

• The fetchUrlTool tool and
• the searchBloomreach tool

@Component
@Order(10)
public final class ExternalUrlToolPackage implements ClientToolPackage {
    public ExternalUrlToolPackage() {
        log.info("ExternalUrlToolPackage is enabled");
    }

    @Override
    public boolean isEnabled() {
        return true;
    }

    @Override
    public String getAdvice() {
        return """
                If the user asks for the content of an external url, use the 'fetchUrlTool' tool passing the url as an argument. The
                tool will provide you with the html content of the webpage, then make a text-based summary of it and present it to the user. 
                Use markdown if you think it can help preserving the formatting you see in the webpage's html.
                
                If the user is requesting to search the Bloomreach documentation for some information, convert this information into a space 
                separated list of keywords and use tool 'searchBloomreach' passing that information as an argument. Then process the html result and
                present the results to the user.
                """;
    }
    @Override
    public Object getTools() {
        return new Object() {

            @Tool(name = "fetchUrlTool", description = "Fetch content of a webpage")
            Document fetchUrlTool(
                    @ToolParam(description = "The url of the webpage to fetch content for") String url,
                    ToolContext toolContext) {

                final HttpClient client = HttpClient.newHttpClient();
                try {
                    final String htmlResponse = client.send(
                             HttpRequest.newBuilder()
                                     .uri(URI.create(url))
                                     .GET()
                                     .build(),
                             HttpResponse.BodyHandlers.ofString())
                       .body();

                    return Document.builder()
                            .id(url)
                            //TODO Html content is too large for the AI, use jsoup
                            .text(htmlResponse)
                            .metadata("url", url)
                            .build();
                } catch (IOException e) {
                    return error("Failed to fetch content of webpage", e);
                } catch (InterruptedException e) {
                    return error("Timed out while fetching content of webpage", e);
                }
            }

            @Tool(name = "searchBloomreach", description = "Perform a search for some keywords against the Bloomreach documentation site")
            //TODO Should return List<Document> instead of asking the AI to process raw html
            Document searchBloomreach(
                    @ToolParam(description = "Keywords that must be used for the search") String keywords,
                    ToolContext toolContext) {

                final String url = String.format("https://xmdocumentation.bloomreach.com/librarysearch?query=%s",
                        URLEncoder.encode(keywords, StandardCharsets.UTF_8));
             
                //TODO Process the html result and return List<Document>
                return fetchUrlTool(url, toolContext);
            }

            private static Document error(final String friendlyMessage, final Exception e) {
                log.error(friendlyMessage, e);
                return Document.builder()
                        .text(String.format("Error, %s: %s",
                                friendlyMessage, e.getMessage()))
                        .build();
            }
        };
    }
}