Skip to content

Commit

Permalink
Refactor chat history & system prompt handling & Store RAG sources in…
Browse files Browse the repository at this point in the history
… history

Signed-off-by: Florian Hotze <[email protected]>
  • Loading branch information
florian-h05 committed Feb 7, 2025
1 parent fcce1da commit 13b8cc8
Show file tree
Hide file tree
Showing 13 changed files with 231 additions and 376 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ private MetadataKeys() {}
*/
public static final String KNOWLEDGE_ID = "knowledge_id";

/** The checksum of the knowledge source. */
/** The checksum of the knowledge content. */
public static final String CHECKSUM = "checksum";

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand All @@ -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:
Expand All @@ -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.
Expand Down
142 changes: 142 additions & 0 deletions src/main/java/com/github/llamara/ai/internal/chat/ChatModel.java
Original file line number Diff line number Diff line change
@@ -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 -> {});
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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) {}
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 =
Expand All @@ -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);
}
}
Expand Down

This file was deleted.

Loading

0 comments on commit 13b8cc8

Please sign in to comment.