ChatBot with Embabel: A complete walkthrough

@Bean
LuceneSearchOperations luceneSearchOperations(ModelProvider modelProvider) {
    return LuceneSearchOperations
            .withName("sources")
            .withEmbeddingService(modelProvider.getEmbeddingService(DefaultModelSelectionCriteria.INSTANCE))
            .withChunkerConfig(new ContentChunker.Config(800,100, 100))
            .withChunkTransformer(AddTitlesChunkTransformer.INSTANCE)
            .withIndexPath(Path.of("./.lucene-index"))
            .buildAndLoadChunks();
}

@Bean
Chatbot chatbot(AgentPlatform agentPlatform) {
    return AgentProcessChatbot.utilityFromPlatform(agentPlatform, new Verbosity().showPrompts());
}

@EmbabelComponent
public class ChatActions {
    private static final Logger logger = LoggerFactory.getLogger(ChatActions.class);

    private final ToolishRag toolishRag;

    public ChatActions(SearchOperations searchOperations) {
        this.toolishRag = new ToolishRag(
                "sources",
                "sources for answering user questions",
                searchOperations);
    }

    @Action(canRerun = true, trigger = UserMessage.class)
    public void respond(Conversation conversation, ActionContext context) {
        String userName = SecurityContextHolder.getContext().getAuthentication().getName();
        var lastUserMessage = conversation.lastMessageIfBeFromUser();

        if (lastUserMessage != null) {
            logger.info("Received user message as last message: {}", lastUserMessage.getContent());
            var assistantMessage = context.ai()
                    .withLlmByRole(CHEAPEST.name())
                    .withReferences(toolishRag)
                    .withSystemPrompt("You are a helpful assistant. Answer questions concisely. Always address the current user by their name: " + userName + ".")
                    .respond(conversation.getMessages());
            context.sendMessage(conversation.addMessage(assistantMessage));
        } else {
            logger.info("Received non-user message");
        }
    }
}

@PostMapping(value = "/init", consumes = "application/json")
public InitSessionResponse initialiseSession(@RequestBody InitSessionRequest request, Authentication authentication) {
    User user = getUser(authentication);
    logger.info("Received request to initialise session for user: {}", user.getId());

    ChatSession chatSession = createOrFetchSession(request.conversationId(), user);

    return new InitSessionResponse(chatSession.getConversation().getId());
}

@NotNull
private ChatSession createOrFetchSession(String conversationId, User user) {
    ChatSession chatSession;
    if (conversationId == null || conversationId.isEmpty()) {
        logger.info("Creating new conversation for user: {}", user.getDisplayName());
        chatSession = chatbot.createSession(user, new SseEmitterOutputChannel(), null);
    } else {
        logger.info("Fetching conversation for ID: {} for user: {}", conversationId, user.getDisplayName());
        chatSession = chatbot.findSession(conversationId);
        if (chatSession == null) {
            throw new IllegalArgumentException("Conversation not found for ID: " + conversationId);
        }
    }
    return chatSession;
}

@GetMapping(value = "/stream/{processId}")
public SseEmitter streamMessages(@PathVariable(name = "processId") String processId, Authentication authentication) {
    logger.info("Starting message streaming for process ID: {}", processId);

    var emitter = new SseEmitter(Long.MAX_VALUE);

    processEmitters.computeIfAbsent(processId, id -> new CopyOnWriteArrayList<>())
            .add(emitter);

    emitter.onCompletion(() -> processEmitters.get(processId).remove(emitter));
    emitter.onError(throwable -> processEmitters.get(processId).remove(emitter));
    emitter.onTimeout(() -> processEmitters.get(processId).remove(emitter));

    User user = getUser(authentication);

    var outputChannel = createOrFetchSession(processId, user).getOutputChannel();
    if (outputChannel instanceof SseEmitterOutputChannel sseEmitterOutputChannel) {
        sseEmitterOutputChannel.setEmitter(emitter, processId);
    } else {
        throw new IllegalStateException("Output channel is not a SseEmitterOutputChannel");
    }

    return emitter;
}

/**
 * Implementation of {@link OutputChannel} that uses Spring's {@link SseEmitter} for Server-Sent Events (SSE).
 * <p>
 * When an emitter is set, it sends a {@code "Connected"} event indicating that the channel is ready
 * to receive events. On reconnecting a client, the OutputChannel can receive a new emitter.
 * </p>
 * <p>
 * Individual {@link OutputChannelEvent} instances are sent sequentially to the client.
 * </p>
 */
public class SseEmitterOutputChannel implements OutputChannel {
    private static final Logger logger = LoggerFactory.getLogger(SseEmitterOutputChannel.class);

    private SseEmitter emitter;

    /**
     * Sets the SSE emitter to use for sending events. The output channel can receive new  emitters.
     *
     * @param emitter The active Sse emitter for sending the events.
     */
    public void setEmitter(SseEmitter emitter, String processId) {
        logger.info("Setting emitter for process: {}", processId);
        this.emitter = emitter;

        // Create the connected event to send to the client
        var sseEvent = SseEmitter.event()
                .id(processId)
                .name("Connected")
                .data(new SseEmitterOutputChannel.ConnectedOutputChannelEvent(processId))
                .build();

        // Send the connected event to the client to let it know that the channel is ready
        try {
            emitter.send(sseEvent);
        } catch (IOException e) {
            throw new OutputChannelRuntimeException("Problem sending the sse stream connected event.", e);
        }
    }

    @Override
    public void send(@NotNull OutputChannelEvent event) {

        if (emitter == null) {
            logger.warn("No emitter set, dropping event: {}", event);
            return;
        }

        // Create the event to send to the client
        var sseEvent = SseEmitter.event()
                .id(event.getProcessId())
                .name(event.getClass().getSimpleName())
                .data(event)
                .build();

        // Send the event to the client
        try {
            emitter.send(sseEvent);
        } catch (IOException e) {
            logger.error("Error sending event to client for processId: {}", event.getProcessId(), e);
            throw new OutputChannelRuntimeException("Error sending event to client", e);
        }
    }

    /**
     * Event for notifying the client that the OutputChannel is ready to receive events.
     */
    public static class ConnectedOutputChannelEvent implements OutputChannelEvent {

        private final String processId;
        private final String message;

        public ConnectedOutputChannelEvent(String processId) {
            this.processId = processId;
            this.message = "Connected OutputChannel SSE Emitter";
        }

        @NotNull
        @Override
        public String getProcessId() {
            return this.processId;
        }

        @NotNull
        public String getMessage() {
            return this.message;
        }
    }
}

@PostMapping(value = "/message", consumes = "application/json")
public Response chat(@RequestBody Request request) {
    logger.info("Received message: {}", request.message());

    var conversationId = request.conversationId();

    // Load the ChatSession using the conversationId or create a new one
    ChatSession chatSession = chatbot.findSession(conversationId);
    if (chatSession == null) {
        throw new IllegalArgumentException("Conversation not found for ID: " + conversationId);
    }

    // Call the agent with the user message
    chatSession.onUserMessage(new UserMessage(request.message()));

    return new Response("You should receive a response soon", chatSession.getProcessId());
}