[FEATURE] Retrieve Documents when using AiService

[KaisNeffati]

Context

#660

Change

Checklist

Before submitting this PR, please check the following points:

• I have added unit and integration tests for my change
• All unit and integration tests in the module I have added/changed are green
• All unit and integration tests in the core and main modules are green
• I have added/updated the documentation
• I have added an example in the examples repo (only for "big" features)
• I have added my new module in the BOM (only when a new module is added)

Checklist for adding new embedding store integration

• I have added a {NameOfIntegration}EmbeddingStoreIT that extends from either EmbeddingStoreIT or EmbeddingStoreWithFilteringIT

[langchain4j]

@KaisNeffati thank you a lot for your contribution!

[langchain4j]

minor: it would be nice to do this validation and fail on AI Service creation (eager), not on usage (lazy).

[langchain4j]

query transformer is not necessary for this test, please remove it to simplify the test

[langchain4j]

Wondering: should AI Service creation fail also when return type is WIthSources, but no retriever/augmentor is set up?

[KaisNeffati]

Absolutely! I agree, adopting a straightforward solution like that would indeed be beneficial. Otherwise, we might encounter numerous complaints from users about the WithSources object lacking contents or other augement information

[langchain4j]

query transformer is not necessary for this test, please remove it to simplify the test

[langchain4j]

please add a test for the scenario when structured output is required: WithSources<Pojo>

[langchain4j]

After a second thought, I think we should make this class a bit more generic, to accomodate for other useful information, such as:

• tools used during the call
• token usage
• etc

(Please do not add these in this PR.)

Any ideas for the good name for the class that can hold that?

[KaisNeffati]

ResponseMessage or ResponsePrompt

Well when we talk about framework like Spring or the javax package, they provide powerful tools for handling HTTP requests and responses. One key concept is how they handle responses.

In Spring, for example, ResponseEntity is a class that encapsulates an HTTP response, including the response body and additional metadata such as headers and status codes. It's a neat way to package up all the necessary information about a response in one object.

Similarly, in the javax package, there's the ResponseBuilder class. This class allows you to construct an HTTP response step by step, setting things like the status code, headers, and entity (the response body). It's a flexible and convenient way to craft custom responses

Now, let's bring this analogy back to our prompt usecase. Just as ResponseEntity and ResponseBuilder provide a structured way to handle HTTP responses. ResponseMessage class aims to provide a structured way to handle prompts response

By using ResponseMessage (or ResponsePrompt), we can create responses that not only contain the text of the prompt itself but also include additional contextual information, just like how ResponseEntity includes metadata with the response body. This makes our response handling more organized, easier to manage, and more adaptable to different scenarios

ResponseMessage responseMessage = ResponseMessage.create("How can I assist you today?")
    .contents(...)
    .tools(...)
    .tokens(...)
    .build();

I can include this in a seperated PR if needed

[langchain4j]

I am not sure about these names, they are not really responses, and not messages/prompts.

These object can contain much more:

• Answer/Response itself (String or CustomPOJO or smth else)
• Retrieval:
  • List<Content> retrievedContents
  • Used queries (not sure if needed)
  • Used retrievers (not sure if needed)
• LLM calls:
  • TokenUsage (for all calls combined and/or for each call separately)
  • FinishReason (of the last call)
  • Model parameters: modelName/temperature/etc?
  • List of all calls with requests and responses (not sure if needed)?
• Tool executions:
  • Executed tools and their inputs/outputs

I am having a hard time to find a good name right now, maybe it is an indicator of a bad design, or maybe I need to go and rest a bit 😆

Have a nice weekend!

[KaisNeffati]

Taking a breather and diving into some meditation is the perfect recipe for this naming problem 😙

As for me, I'm rolling with 'WithSources' for now. But if you've got more brainwaves brewing, feel free to open another issue, and we'll keep this brainstorm party going! Yep, it's definitely more than just a response 🤔

[langchain4j]

This is a breaking change, is there a way to avoid this?

[KaisNeffati]

The purpose of this is to retrieve the contents used for augmentation.

I have two solutions in mind:

1. First solution: Add another method in the RetrievalAugmentor which returns another object containing the UserMessage and the list of contents.
2. Second solution: Add the list of contents (List<TextSegments>) in the UserMessage class.

I noticed another list of contents in the UserMessage, but it's only used to specify the type of content used for augmentation.

What do you suggest to make retrieving the UserMessage with the list of content possible in the DefaultAiServices class?

[langchain4j]

I assume that in the long term, regardless of this change, we will have to make internals of UserMessage more fine-grained, e.g. having user message and retrieved contents separate (solution 2). I foresee that LLM APIs in the future will separate user message from supporting (retrieved) content to be able to adjust/fine-tune their LLMs better. Some LLM providers already do this (see documents parameter). But this is pretty big change, and introduces more complexity, so I would better postpone it untill it will be really required.

I guess for now we can go with the first solution, something like this?

public interface RetrievalAugmentor {

    @Deprecated // let's deprecate this
    UserMessage augment(UserMessage userMessage, Metadata metadata); // existing API

    default AugmentationResult augment(AugmentationRequest augmentationRequest) { // new API
        UserMessage augmented = augment(augmentationRequest.chatMessage, augmentationRequest.metadata);
        return new AugmentationResult(augmented, null);
    }
}

// introduce request/result classes to make it more extendable in the future
public class AugmentationRequest {

    ChatMessage chatMessage; // instead of UserMessage, this will be useful for issue #793
    Metadata metadata;
}

public class AugmentationResult {

    ChatMessage augmentedChatMessage;
    List<Content> retrievedContents;
}

This feature will not work for existing custom implementations of RetrievalAugmentor, but I guess it is fine as there are most probably not many such users and this class is explicitly marked @Experimental.

Most users are probably using DefaultRetrievalAugmentor, where we will have a proper implementation of the new method.

WDYT?

[KaisNeffati]

I'll include the ChatMessage in a separate PR since it requires numerous modifications

[langchain4j]

@KaisNeffati what is the change exactly? Is it breaking?

[KaisNeffati]

To utilize ChatMessage within both AugmentationRequest and AugmentationResult, the interface should be adjusted as follows:

    @Deprecated
    UserMessage augment(UserMessage userMessage, Metadata metadata);

    UserMessage augment(ChatMessage chatMessage, Metadata metadata);

    /**
     * Augments the provided {@link AugmentationRequest} with retrieved content.
     *
     * @param augmentationRequest The {@link AugmentationRequest} containing the user message and metadata.
     * @return The {@link AugmentationResult} containing the augmented user message.
     */
    default AugmentationResult augment(AugmentationRequest augmentationRequest) { // new API
        UserMessage augmented = augment(augmentationRequest.getChatMessage(), augmentationRequest.getMetadata());
        return AugmentationResult.builder()
                .augmentedUserMessage(augmented)
                .build();
    }

This adjustment necessitates modifying the contentInjector to accept:

UserMessage inject(List<Content> contents, ChatMessage chatMessage);
etc...

So for this PR i kept using the UserMessage , And i'll include the use of ChatMessage in another PR

[langchain4j]

Suggested change

	private AugmentedMessage augmentedMessage; // Wrapper for the augmentation details.
	private List<Content> retrievedContents;

WDYT?

[langchain4j]

Hi @KaisNeffati, do you mind if I will do and push some changes?

[KaisNeffati]

Feel free to do it @langchain4j