Add prompt-formatter functionality

[andreadimaio]

Hi, I would like to use this topic to share with you an idea that I have in mind and understand if this is something that could be used for one or more LLM providers present in quarkus-langchain4j.

Why introduce the prompt-formatter?

I mainly work with the watsonx provider, which allows the developer to use different models. Some of these models (especially those ending with -instruct) can be enriched with special tags that allow to get better results from the LLM, but this is not done automatically by watsonx.

To better understand what I'm saying, here's an example.

@RegisterAiService
public interface LLMService {

    @SystemMessage("""
        You are a calculator and you calculate the sum of two numbers given as input.
        """)
    @UserMessage("""
        Input:
        {firstNumber} + {firstNumber}
        Output:
        """)
    public String calculator(int firstNumber, int secondNumber);
}

Suppose that I want to use this prompt with three different models: llama-3-405b-instruct/mistral-large/granite-13b-instruct-v2

To get the best response, the prompt should be written like this:

llama-3-405b-instruct:

<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a calculator and you calculate the sum of two numbers given as input.<|eot_id|>
<|start_header_id|>user<|end_header_id|>
Input:
{firstNumber} + {firstNumber}
Output:<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>

mistral-large:

<s>[INST]You are a calculator and you calculate the sum of two numbers given as input.[/INST]</s>
[INST]Input:
{firstNumber} + {firstNumber}
Output:[/INST]

granite-13b-instruct-v2:

<|system|>
You are a calculator and you calculate the sum of two numbers given as input.
<|user|>
Input:
{firstNumber} + {firstNumber}
Output:
<|assistant|>

So what I normally do to get better results from the LLM is to add the tags I need based on the model, but this becomes a problem if I then want to change it, because I have to re-edit the whole prompt. The idea is to have a class that does the work automatically.

PromptFormatter interface

I already have a draft implementation (the examples above were automatically generated) which consists of having a PromptFormatter interface and a class for each model like LLamaPromptFormatter, MistralPromptFormatter and GranitePromptFormatter. The idea is to create the PromptFormatter class based on the model name during the build phase, so that this class is used during the runtime phase to enrich the prompt with the correct tags.

/**
 * The PromptFormatter interface defines the structure for handling and converting different types of {@link ChatMessage}
 * objects
 * into a specific string format for prompt generation.
 */
public interface PromptFormatter {

    /**
     * Defines the string used to join multiple {@link ChatMessage} objects when constructing the prompt.
     *
     * @return the string used to join messages in the prompt.
     */
    public default String joiner() {
        return "\n";
    }

    /**
     * Defines a start tag that wraps the entire prompt.
     *
     * @return the start tag for the prompt.
     */
    public default String start() {
        return "";
    }

    /**
     * Defines an end tag that wraps the entire prompt.
     *
     * @return the end tag for the prompt.
     */
    public default String end() {
        return "";
    }

    /**
     * Returns the tag used to mark {@link SystemMessage} within the prompt.
     *
     * @return the tag representing a system message.
     */
    String system();

    /**
     * Returns the tag used to mark {@link UserMessage} within the prompt.
     *
     * @return the tag representing a user message.
     */
    String user();

    /**
     * Returns the tag used to mark {@link AiMessage} within the prompt.
     *
     * @return the tag representing an assistant message.
     */
    String assistant();

    /**
     * Returns the tag used to mark {@link ToolExecutionResultMessage} within the prompt.
     *
     * @return the tag representing a tool execution message.
     */
    String tool();

    /**
     * Defines how to close a tag based on the message type.
     *
     * @param type the {@link ChatMessageType} for which the closing tag is being requested.
     * @return the closing tag for the specified message type.
     */
    String endOf(ChatMessageType type);

    /**
     * Defines how to close a tag based on the message.
     *
     * @param message the {@link ChatMessage} for which the closing tag is being requested.
     * @return the closing tag for the specified message type.
     */
    default String endOf(ChatMessage message) {
        return endOf(message.type());
    }

    /**
     * Returns the tag associated with a specific {@link ChatMessageType}.
     *
     * @param type the {@link ChatMessageType} for which the tag is being requested.
     * @return the tag for the specified message type.
     */
    default String tagOf(ChatMessageType type) {
        return switch (type) {
            case AI -> assistant();
            case SYSTEM -> system();
            case TOOL_EXECUTION_RESULT -> tool();
            case USER -> user();
        };
    }

    /**
     * Returns the tag associated with a specific {@link ChatMessage}.
     *
     * @param message the {@link ChatMessage} for which the tag is being requested.
     * @return the tag for the specified message.
     */
    default String tagOf(ChatMessage message) {
        return tagOf(message.type());
    }

    /**
     * Returns a list of all relevant tags used in the prompt.
     *
     * @return a list of all relevant tags used in the prompt.
     */
    default List<String> tokens() {
        return Stream.of(start(), end(), system(), user(), assistant(), tool())
                .map(String::trim)
                .filter(not(String::isBlank))
                .toList();
    }
}

chat-model.prompt-formatter property

The prompt-formatter feature can be enable or disabled with the chat-model.prompt-formatter property.
The idea is to have two different values:

• enabled: The feature is explicitly enabled, and tags are added automatically.
• disabled: The feature is explicitly disabled, and no tags are added.

The enabled values ​​need checks during build to throw an exception if misconfigured, for example, if the value is enabled and a prompt contains one of the special tags, an exception must be thrown.

Final consideration

For now, this is just a draft, and I'm sure there's more to add or change to get something that works. But I wanted to share it with you to get your opinion on the idea and to understand if this feature makes sense only for watsonx or if it can be added to other providers that allow multiple models (like ollama or hugging-face).

[andreadimaio]

Looking at the API of ollama and hugging-face it seems that they support "chat completion". This means that the two providers should add the tags automatically.

[andreadimaio]

This means that if we decide to implement this new feature, it will only be related to the watsonx provider.

[geoand]

Thanks for starting this discussion!

I am pretty sure that this is something @cescoffier wanted as well, so let's wait until he is back from PTO

[andreadimaio]

Here is a draft implementation (for watsonx only) of the idea I described above.