Spring-ai [Messages] Consolidate relationship between Generation and AssistantMessage

Expected Behavior

The Assistant role is used for messages sent by AI services in response to user prompts. Considering that implementations of the Message interface exist for each role, I would expect the AssistantMessage implementation to be used as part of a ChatResponse to wrap the reply from the AI service.

Current Behavior

The response from an AI service is wrapped in a Generation object, which is also a Message, but not an AssistantMessage. I can see the need for this special object as it includes additional metadata we want and need as part of a response from an AI service. At the same time, if the answer needs to be used by the application as part of subsequent prompts to maintain the conversation history (something that will probably be handled by Spring AI in the future), an explicit operation is needed to build an AssistantMessage object with information extracted from the Generation object.

The Generation object is returned as part of a ChatResponse with message type "Assistant". What's the relationship with AssistantMessage and can such relationship be made more explicit? For example, could Generation be a specialisation of AssistantMessage? Could Generation expose a method to return an AssistantMessage?

Comment From: ThomasVitale

Looks like the changes in https://github.com/spring-projects/spring-ai/pull/249/files#diff-b6731b45050291dc99d46a82622c92ed1e6164559519328153240bacea67139b would address this issue and consolidate the relationship between Generation and AssistantMessage.

Comment From: markpollack

I think the changes address your concern. Also, there is a branched named "options-ftw" that will better handle the metadata,

public class Generation implements ModelResult<AssistantMessage> {

    private AssistantMessage assistantMessage;

    private ChatGenerationMetadata chatGenerationMetadata;

Please reopen this issue if you think something is amiss.