Skip to content

Commit

Permalink
Merge branch 'refs/heads/chat-client-refactoring'
Browse files Browse the repository at this point in the history
  • Loading branch information
@jschm42
jschm42 committed Jun 20, 2024
2 parents 019ba49 + 4693360 commit de2d72d
Show file tree
Hide file tree
Showing 18 changed files with 259 additions and 155 deletions.
4 changes: 2 additions & 2 deletions backend/src/main/java/com/talkforgeai/backend/assistant/dto/AssistantDto.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ public record AssistantDto(

public enum MemoryType {
NONE,
ASSISTANT,
GLOBAL
AI_DECIDES,
HISTORY
}
}
91 changes: 46 additions & 45 deletions backend/src/main/java/com/talkforgeai/backend/assistant/service/AssistantSpringService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,14 @@
import com.talkforgeai.backend.assistant.dto.ThreadTitleGenerationRequestDto;
import com.talkforgeai.backend.assistant.dto.ThreadTitleUpdateRequestDto;
import com.talkforgeai.backend.assistant.exception.AssistentException;
import com.talkforgeai.backend.assistant.functions.ContextStorageFunction;
import com.talkforgeai.backend.assistant.functions.ContextStorageFunction.Request;
import com.talkforgeai.backend.assistant.functions.ContextStorageFunction.Response;
import com.talkforgeai.backend.assistant.functions.ContextTool;
import com.talkforgeai.backend.assistant.functions.FunctionContext;
import com.talkforgeai.backend.assistant.repository.AssistantRepository;
import com.talkforgeai.backend.assistant.repository.MessageRepository;
import com.talkforgeai.backend.assistant.repository.ThreadRepository;
import com.talkforgeai.backend.memory.dto.DocumentWithoutEmbeddings;
import com.talkforgeai.backend.memory.functions.MemoryContextStorageFunction;
import com.talkforgeai.backend.memory.functions.MemoryContextStorageFunction.Request;
import com.talkforgeai.backend.memory.functions.MemoryContextStorageFunction.Response;
import com.talkforgeai.backend.memory.functions.MemoryFunctionContext;
import com.talkforgeai.backend.memory.service.MemoryService;
import com.talkforgeai.backend.storage.FileStorageService;
import com.talkforgeai.backend.transformers.MessageProcessor;
Expand Down Expand Up @@ -79,12 +78,14 @@
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi.ChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.data.domain.PageRequest;
Expand Down Expand Up @@ -157,7 +158,7 @@ public AssistantSpringService(
return runIdMono;
}

private static @NotNull List<Message> getFinalPromptMessageList(String message,
private static @NotNull List<Message> getFinalPromptMessageList(
List<MessageDto> pastMessagesList, AssistantDto assistantDto,
List<DocumentWithoutEmbeddings> memoryResultsList) {
List<Message> promptMessageList = pastMessagesList.stream()
Expand Down Expand Up @@ -194,9 +195,6 @@ public AssistantSpringService(
result -> memoryMessage.append(result.content()).append("\n"));
memoryMessage.append("\nUser message:\n");
}

finalPromptMessageList.add(
new UserMessage(memoryMessage.append(message).toString()));
return finalPromptMessageList;
}

Expand Down Expand Up @@ -258,11 +256,11 @@ public void cancelStream(String threadId, String runId) {
}

public Flux<ServerSentEvent<String>> streamRunConversation(String assistantId, String threadId,
String message) {
String userMessage) {

final String runId = UniqueIdUtil.generateRunId();

Mono<Object> saveUserMessageMono = getSaveUserMessageMono(assistantId, threadId, message);
Mono<Object> saveUserMessageMono = getSaveUserMessageMono(assistantId, threadId, userMessage);
Mono<AssistantDto> assistantEntityMono = getAssistantEntityMono(assistantId);
Mono<List<MessageDto>> pastMessages = getPastMessagesMono(threadId);
Mono<InitInfos> initInfosMono = getInitInfosMono(assistantEntityMono, pastMessages);
Expand All @@ -275,7 +273,7 @@ public Flux<ServerSentEvent<String>> streamRunConversation(String assistantId, S
.flux()
.flatMap(initInfos -> {
List<DocumentWithoutEmbeddings> memorySearchResults = getMemorySearchResults(
initInfos.assistantDto, message);
initInfos.assistantDto.id(), initInfos.assistantDto.memory(), userMessage);

return Flux.just(
new PreparedInfos(initInfos.assistantDto(), initInfos.pastMessages(),
Expand All @@ -286,36 +284,40 @@ public Flux<ServerSentEvent<String>> streamRunConversation(String assistantId, S
List<MessageDto> pastMessagesList = preparedInfos.pastMessages();
List<DocumentWithoutEmbeddings> memoryResultsList = preparedInfos.memoryResults();

List<Message> finalPromptMessageList = getFinalPromptMessageList(message,
List<Message> finalPromptMessageList = getFinalPromptMessageList(
pastMessagesList,
assistantDto, memoryResultsList);

FunctionCallbackWrapper<Request, Response> memoryFunctionCallback = getMemoryFunctionCallback(
assistantId, runId, assistantDto);
List<FunctionCallback> functionCallbacks = new ArrayList<>();
FunctionCallbackWrapper<Request, Response> memoryFunctionCallback
= getMemoryFunctionCallback(assistantId, assistantDto.name(), runId);

if (assistantDto.memory() == MemoryType.AI_DECIDES) {
functionCallbacks.add(memoryFunctionCallback);
}

ChatOptions promptOptions = universalChatService.getPromptOptions(assistantDto,
List.of(memoryFunctionCallback));
functionCallbacks);

LOGGER.debug("Starting stream with prompt: {}", finalPromptMessageList);
LOGGER.debug("Prompt Options: {}",
universalChatService.printPromptOptions(assistantDto.system(),
promptOptions));

Prompt prompt = new Prompt(finalPromptMessageList, promptOptions);

return universalChatService.stream(assistantDto.system(), prompt);
return universalChatService.stream(assistantDto,
finalPromptMessageList, userMessage, promptOptions);
})
.doOnCancel(() -> {
LOGGER.debug("doOnCancel. message={}", assistantMessageContent);
LOGGER.debug("doOnCancel. userMessage={}", assistantMessageContent);
})
.mapNotNull(chatResponse -> mapChatResponse(chatResponse, assistantMessageContent))
.doOnSubscribe(subscription -> {
LOGGER.debug("doOnSubscribe. message={}", assistantMessageContent);
LOGGER.debug("doOnSubscribe. userMessage={}", assistantMessageContent);

activeStreams.put(runId, subscription);
})
.doOnComplete(() -> {
LOGGER.trace("doOnComplete. message={}", assistantMessageContent);
LOGGER.trace("doOnComplete. userMessage={}", assistantMessageContent);

Mono.fromRunnable(
() -> saveNewMessage(assistantId, threadId, MessageType.ASSISTANT,
Expand Down Expand Up @@ -356,18 +358,18 @@ public Flux<ServerSentEvent<String>> streamRunConversation(String assistantId, S
}

private FunctionCallbackWrapper<Request, Response> getMemoryFunctionCallback(
String assistantId, String runId, AssistantDto assistantDto) {
String assistantId, String assistantName, String runId) {
return FunctionCallbackWrapper.builder(
new ContextStorageFunction(memoryService,
new FunctionContext(LlmSystem.OPENAI,
new MemoryContextStorageFunction(memoryService,
new MemoryFunctionContext(LlmSystem.OPENAI,
OpenAiEmbeddingProperties.DEFAULT_EMBEDDING_MODEL,
assistantId,
assistantDto.name(),
assistantName,
runId
)))
.withDescription(
"Store relevant information in the vector database for later retrieval.")
.withName(ContextTool.MEMORY_STORE.getFunctionBeanName())
.withName(MemoryContextStorageFunction.NAME)
.build();
}

Expand All @@ -394,31 +396,32 @@ private FunctionCallbackWrapper<Request, Response> getMemoryFunctionCallback(
.subscribeOn(Schedulers.boundedElastic());
}

private @NotNull List<DocumentWithoutEmbeddings> getMemorySearchResults(AssistantDto assistantDto,
private @NotNull List<DocumentWithoutEmbeddings> getMemorySearchResults(String assistantId,
MemoryType memoryType,
String message) {

if (assistantDto.memory() == MemoryType.NONE) {
if (memoryType == MemoryType.NONE) {
return List.of();
}

LOGGER.info("Searching memory for message: {}", message);
List<DocumentWithoutEmbeddings> searchResults = memoryService.search(
SearchRequest.query(message).withSimilarityThreshold(0.75f));

if (assistantDto.memory() == MemoryType.ASSISTANT) {
List<DocumentWithoutEmbeddings> filteredMemory = searchResults.stream()
.filter(m -> m.assistantId() != null && m.assistantId().equals(assistantDto.id()))
.toList();
FilterExpressionBuilder expressionBuilder = new FilterExpressionBuilder();
Expression assistantExpression = expressionBuilder.eq("assistantId", assistantId).build();

LOGGER.debug("Memory search results for assistant '{}': {}", assistantDto.id(),
filteredMemory);
List<DocumentWithoutEmbeddings> searchResults = memoryService.search(
SearchRequest.query(message)
.withFilterExpression(assistantExpression)
.withSimilarityThreshold(0.75f));

return filteredMemory;
}
List<DocumentWithoutEmbeddings> filteredMemory = searchResults.stream()
.filter(m -> m.assistantId() != null && m.assistantId().equals(assistantId))
.toList();

LOGGER.debug("Memory search results: {}", searchResults);
LOGGER.debug("Memory search results for assistant '{}': {}", assistantId,
filteredMemory);

return searchResults;
return filteredMemory;
}

private ServerSentEvent<String> createResponseSseEvent(ChatResponse chatResponse) {
Expand Down Expand Up @@ -521,10 +524,8 @@ public ThreadTitleDto generateThreadTitle(String threadId,
.withMaxTokens(256)
.build();

Prompt titlePrompt = new Prompt(new UserMessage(content), options);

try {
ChatResponse titleResponse = universalChatService.call(LlmSystem.OPENAI, titlePrompt);
ChatResponse titleResponse = universalChatService.call(LlmSystem.OPENAI, content, options);

String generatedTitle = titleResponse.getResult().getOutput().getContent();

Expand Down
85 changes: 57 additions & 28 deletions backend/src/main/java/com/talkforgeai/backend/assistant/service/UniversalChatService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,19 +17,24 @@
package com.talkforgeai.backend.assistant.service;

import com.talkforgeai.backend.assistant.dto.AssistantDto;
import com.talkforgeai.backend.assistant.dto.AssistantDto.MemoryType;
import com.talkforgeai.backend.assistant.dto.LlmSystem;
import com.talkforgeai.backend.assistant.exception.AssistentException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.anthropic.AnthropicChatModel;
import org.springframework.ai.anthropic.AnthropicChatOptions;
import org.springframework.ai.anthropic.api.AnthropicApi;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.RequestResponseAdvisor;
import org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.mistralai.MistralAiChatModel;
import org.springframework.ai.mistralai.MistralAiChatOptions;
import org.springframework.ai.mistralai.api.MistralAiApi;
Expand All @@ -38,6 +43,7 @@
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;
Expand All @@ -46,6 +52,8 @@
@Service
public class UniversalChatService {

private final int DEFAULT_CHAT_MEMORY_RESPONSE_SIZE = 5;

@Qualifier("openAiRestClient")
private final RestClient openAiRestClient;
private final OpenAiChatModel openAiChatModel;
Expand All @@ -56,16 +64,22 @@ public class UniversalChatService {
@Qualifier("ollamaAiRestClient")
private final RestClient ollamaAiRestClient;

@Qualifier("dbVectorStore")
private final VectorStore dbVectorStore;

public UniversalChatService(RestClient openAiRestClient,
OpenAiChatModel openAiChatModel,
MistralAiChatModel mistralAiChatModel, AnthropicChatModel anthropicChatModel,
OllamaChatModel ollamaChatModel, RestClient ollamaAiRestClient) {
OllamaChatModel ollamaChatModel,
RestClient ollamaAiRestClient, VectorStore dbVectorStore) {
this.openAiRestClient = openAiRestClient;
this.openAiChatModel = openAiChatModel;
this.mistralAiChatModel = mistralAiChatModel;
this.anthropicChatModel = anthropicChatModel;
this.ollamaChatModel = ollamaChatModel;
this.ollamaAiRestClient = ollamaAiRestClient;

this.dbVectorStore = dbVectorStore;
}

public ChatOptions getPromptOptions(AssistantDto assistantDto,
Expand Down Expand Up @@ -126,38 +140,53 @@ public String printPromptOptions(LlmSystem system, ChatOptions options) {
return printedOptions.toString();
}

ChatResponse call(LlmSystem system, Prompt prompt) {
return getChatClient(system).call(prompt);
ChatResponse call(LlmSystem system, String prompt, ChatOptions options) {
return getClient(system)
.prompt()
.options(options)
.user(prompt)
.call()
.chatResponse();
}

Flux<ChatResponse> stream(LlmSystem system, Prompt prompt) {
return getStreamingChatClient(system).stream(prompt);
}
Flux<ChatResponse> stream(AssistantDto assistantDto, List<Message> messages,
String userMessage, ChatOptions options) {

List<RequestResponseAdvisor> requestResponseAdvisors = new ArrayList<>();

StreamingChatModel getStreamingChatClient(LlmSystem system) {
return (StreamingChatModel) getClient(system);
if (assistantDto.memory() == MemoryType.HISTORY) {
requestResponseAdvisors.add(getVectorStoreChatMemoryAdvisor(assistantDto.id()));
}

return getClient(assistantDto.system())
.prompt()
.advisors(requestResponseAdvisors)
.options(options)
.messages(messages)
.user(userMessage)
.stream()
.chatResponse();
}

ChatModel getChatClient(LlmSystem system) {
return (ChatModel) getClient(system);
private @NotNull VectorStoreChatMemoryAdvisor getVectorStoreChatMemoryAdvisor(
String converationId) {
return new VectorStoreChatMemoryAdvisor(
dbVectorStore,
converationId,
DEFAULT_CHAT_MEMORY_RESPONSE_SIZE
);
}

private Object getClient(LlmSystem system) {
switch (system) {
case OPENAI -> {
return openAiChatModel;
}
case MISTRAL -> {
return mistralAiChatModel;
}
case OLLAMA -> {
return ollamaChatModel;
}
case ANSTHROPIC -> {
return anthropicChatModel;
}
default -> throw new IllegalStateException("Unexpected system: " + system);
}
private ChatClient getClient(LlmSystem system) {
ChatModel model = switch (system) {
case OPENAI -> openAiChatModel;
case MISTRAL -> mistralAiChatModel;
case OLLAMA -> ollamaChatModel;
case ANSTHROPIC -> anthropicChatModel;
};

return ChatClient.builder(model)
.build();
}

private MistralAiChatOptions getMistralOptions(AssistantDto assistantDto,
Expand Down
12 changes: 12 additions & 0 deletions backend/src/main/java/com/talkforgeai/backend/memory/domain/MemoryDocument.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,11 @@ public class MemoryDocument {

@ManyToOne
private AssistantEntity assistant;

@Column(name = "run_id", length = 50)
private String runId;
@Column(name = "message_type", length = 20)
private String messageType;

public String getId() {
return id;
Expand Down Expand Up @@ -146,4 +149,13 @@ public String getRunId() {
public void setRunId(String runId) {
this.runId = runId;
}

public String getMessageType() {
return messageType;
}

public void setMessageType(String memoryType) {
this.messageType = memoryType;
}

}
Loading

0 comments on commit de2d72d

Please sign in to comment.