Refactor chat history/system prompt handling & Store sources in history

src/main/java/com/github/llamara/ai/internal/chat/aiservice/ChatModelAiService.java → src/main/java/com/github/llamara/ai/internal/chat/AiService.java

@@ -17,14 +17,13 @@
  * limitations under the License.
  * #L%
  */
-package com.github.llamara.ai.internal.chat.aiservice;
+package com.github.llamara.ai.internal.chat;
 
 import java.util.UUID;
 
 import dev.langchain4j.service.MemoryId;
 import dev.langchain4j.service.Result;
 import dev.langchain4j.service.SystemMessage;
-import dev.langchain4j.service.TokenStream;
 import dev.langchain4j.service.UserMessage;
 
 /**
@@ -38,7 +37,7 @@
  *
  * @author Florian Hotze - Initial contribution
  */
-public interface ChatModelAiService {
+public interface AiService {
     String SYSTEM_MESSAGE =
             """
 You are LLAMARA, the Large Language Assistant for Model Augmented Retrieval and Analysis:
@@ -50,14 +49,9 @@ public interface ChatModelAiService {
 """;
 
     @SystemMessage(SYSTEM_MESSAGE)
-    Result<String> chat(@MemoryId UUID sessionId, boolean history, @UserMessage String prompt);
+    Result<String> chat(@MemoryId UUID sessionId, @UserMessage String prompt);
 
-    Result<String> chatWithoutSystemMessage(
-            @MemoryId UUID sessionId, boolean history, @UserMessage String prompt);
-
-    @SystemMessage(SYSTEM_MESSAGE)
-    TokenStream chatAndStreamResponse(
-            @MemoryId UUID sessionId, boolean history, @UserMessage String prompt);
+    Result<String> chatWithoutSystemMessage(@MemoryId UUID sessionId, @UserMessage String prompt);
 
     /**
      * Clean the given text by removing unnecessary noise and formatting it.

src/main/java/com/github/llamara/ai/internal/chat/ChatModel.java

@@ -0,0 +1,142 @@
+/*
+ * #%L
+ * llamara-backend
+ * %%
+ * Copyright (C) 2024 - 2025 Contributors to the LLAMARA project
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.github.llamara.ai.internal.chat;
+
+import com.github.llamara.ai.config.chat.ChatModelConfig;
+import com.github.llamara.ai.internal.MetadataKeys;
+import com.github.llamara.ai.internal.chat.history.ChatHistoryStore;
+import com.github.llamara.ai.internal.chat.history.ChatMessageRecord;
+import com.github.llamara.ai.internal.chat.response.ChatResponseRecord;
+import com.github.llamara.ai.internal.chat.response.RagSourceRecord;
+
+import java.time.Instant;
+import java.util.List;
+import java.util.UUID;
+
+import dev.langchain4j.data.message.ChatMessageType;
+import dev.langchain4j.service.Result;
+
+/**
+ * The {@link ChatModel} provides the interface to chat with the chat models. It takes care of
+ * storing the chat history, if enabled, and applies the system prompt, if enabled.
+ *
+ * @author Florian Hotze - Initial contribution
+ */
+public class ChatModel {
+    private final ChatModelConfig.ModelConfig config;
+    private final AiService aiService;
+    private final ChatHistoryStore historyStore;
+
+    public ChatModel(
+            ChatModelConfig.ModelConfig config,
+            AiService aiService,
+            ChatHistoryStore historyStore) {
+        this.config = config;
+        this.aiService = aiService;
+        this.historyStore = historyStore;
+    }
+
+    /**
+     * Send a prompt to a chat model and get the response.
+     *
+     * @param sessionId the session ID
+     * @param history whether to save the conversation to the chat history
+     * @param prompt the prompt to send to the chat model
+     * @return the response from the chat model
+     */
+    public ChatResponseRecord chat(UUID sessionId, boolean history, String prompt) {
+        if (history) {
+            storePrompt(sessionId, prompt);
+        }
+
+        Result<String> result;
+        if (config.systemPromptEnabled()) {
+            result = aiService.chat(sessionId, prompt);
+        } else {
+            result = aiService.chatWithoutSystemMessage(sessionId, prompt);
+        }
+        ChatResponseRecord response =
+                new ChatResponseRecord(
+                        result.content(), getSourcesFromResult(result, result.content()));
+
+        if (history) {
+            storeResponse(sessionId, response);
+        }
+
+        return response;
+    }
+
+    /**
+     * Get the sources from the result. Filters out sources that were not used by the chat model to
+     * generate the response.
+     *
+     * @param result the {@link Result} from the chat model
+     * @param text the text of the response
+     * @return the sources used by the chat model to generate the response
+     */
+    private List<RagSourceRecord> getSourcesFromResult(Result<?> result, String text) {
+        return result.sources().stream()
+                .map(dev.langchain4j.rag.content.Content::textSegment)
+                .filter(ts -> text.contains(ts.metadata().getString(MetadataKeys.KNOWLEDGE_ID)))
+                .map(
+                        ts ->
+                                new RagSourceRecord(
+                                        ts.metadata().getUUID(MetadataKeys.KNOWLEDGE_ID),
+                                        ts.text()))
+                .toList();
+    }
+
+    /**
+     * Store the prompt in the chat history.
+     *
+     * @param sessionId the session ID
+     * @param prompt the prompt to store
+     */
+    private void storePrompt(UUID sessionId, String prompt) {
+        historyStore
+                .addMessage(
+                        sessionId,
+                        new ChatMessageRecord(
+                                ChatMessageType.USER, prompt, Instant.now(), null, null, null))
+                .subscribe()
+                .with(item -> {}, failure -> {});
+    }
+
+    /**
+     * Store the response in the chat history.
+     *
+     * @param sessionId the session ID
+     * @param response the response to store
+     */
+    private void storeResponse(UUID sessionId, ChatResponseRecord response) {
+        historyStore
+                .addMessage(
+                        sessionId,
+                        new ChatMessageRecord(
+                                ChatMessageType.AI,
+                                response.response(),
+                                Instant.now(),
+                                response.sources(),
+                                config.provider(),
+                                config.model()))
+                .subscribe()
+                .with(item -> {}, failure -> {});
+    }
+}

src/main/java/com/github/llamara/ai/internal/chat/ChatModelContainer.java

@@ -20,26 +20,23 @@
 package com.github.llamara.ai.internal.chat;
 
 import com.github.llamara.ai.config.chat.ChatModelConfig;
-import com.github.llamara.ai.internal.chat.aiservice.ChatModelAiService;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 
 /**
  * Container for a chat model, containing information about the chat model and the {@link
- * ChatModelAiService}s.
+ * ChatModel}.
  *
  * @param uid the UID of the chat model
  * @param label the label of the chat model
  * @param description the description of the chat model
  * @param provider the {@link ChatModelConfig.ChatModelProvider}
- * @param config the {@link ChatModelConfig.ModelConfig}
- * @param service the {@link ChatModelAiService} of the chat model
+ * @param model the {@link ChatModel}
  * @author Florian Hotze - Initial contribution
  */
 public record ChatModelContainer(
         String uid,
         String label,
         String description,
         ChatModelConfig.ChatModelProvider provider,
-        @JsonIgnore ChatModelConfig.ModelConfig config,
-        @JsonIgnore ChatModelAiService service) {}
+        @JsonIgnore ChatModel model) {}

src/main/java/com/github/llamara/ai/internal/chat/ChatModelProvider.java

@@ -20,15 +20,14 @@
 package com.github.llamara.ai.internal.chat;
 
 import com.github.llamara.ai.config.chat.ChatModelConfig;
-import com.github.llamara.ai.internal.chat.aiservice.ChatModelAiService;
 
 import java.util.Collection;
 
 /**
  * The chat model provider provides access to the configured {@link ChatModelContainer}s.
  *
- * <p>It processes the {@link ChatModelConfig} and creates the {@link ChatModelAiService}s to
- * interface with the configured models.
+ * <p>It processes the {@link ChatModelConfig} and creates the {@link AiService}s to interface with
+ * the configured models.
  *
  * @author Florian Hotze - Initial contribution
  */

src/main/java/com/github/llamara/ai/internal/chat/ChatModelProviderImpl.java

@@ -25,9 +25,7 @@
 import com.github.llamara.ai.config.EnvironmentVariables;
 import com.github.llamara.ai.config.chat.ChatModelConfig;
 import com.github.llamara.ai.internal.StartupException;
-import com.github.llamara.ai.internal.chat.aiservice.ChatModelAiService;
 import com.github.llamara.ai.internal.chat.history.ChatHistoryStore;
-import com.github.llamara.ai.internal.chat.history.HistoryInterceptingAiService;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -106,15 +104,15 @@ private void initializeChatModels() {
             ChatLanguageModel clm = produceChatLanguageModel(config);
             StreamingChatLanguageModel sclm = produceStreamingChatLanguageModel(config);
 
-            ChatModelAiService service =
-                    new HistoryInterceptingAiService(
-                            AiServices.builder(ChatModelAiService.class)
+            ChatModel model =
+                    new ChatModel(
+                            config,
+                            AiServices.builder(AiService.class)
                                     .chatLanguageModel(clm)
                                     .streamingChatLanguageModel(sclm)
                                     .chatMemoryProvider(chatMemoryProvider)
                                     .retrievalAugmentor(retrievalAugmentor)
                                     .build(),
-                            config,
                             chatHistoryStore);
 
             ChatModelContainer cm =
@@ -123,8 +121,7 @@ private void initializeChatModels() {
                             config.label().orElse(config.uid()),
                             config.description().orElse(config.provider() + " " + config.model()),
                             config.provider(),
-                            config,
-                            service);
+                            model);
             chatModels.put(config.uid(), cm);
         }
     }