Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add ChatAgent and basic DefaultChatAgent

Browse Source
This commit is contained in:
Mark Pollack
2024-04-18 15:41:28 -04:00
parent e45a175a48
commit a0847cc2cb
9 changed files with 569 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiDefaultChatAgentIT.java Normal file
View File

@@ -0,0 +1,108 @@
package org.springframework.ai.openai.chat.agent;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.agent.DefaultChatAgent;
import org.springframework.ai.chat.agent.PromptContext;
import org.springframework.ai.chat.agent.transformer.QAPromptContextTransformer;
import org.springframework.ai.chat.agent.transformer.VectorStorePromptContextTransformer;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.Resource;
import java.util.List;
@SpringBootTest(classes = OpenAiDefaultChatAgentIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class OpenAiDefaultChatAgentIT {
private final ChatClient chatClient;
private final VectorStore vectorStore;
@Value("classpath:/data/acme/bikes.json")
private Resource bikesResource;
@Autowired
public OpenAiDefaultChatAgentIT(ChatClient chatClient, VectorStore vectorStore) {
this.chatClient = chatClient;
this.vectorStore = vectorStore;
}
@Test
void simpleChat() {
loadData();
var chatAgent = DefaultChatAgent.builder(chatClient)
.withRetrievers(List.of(new VectorStorePromptContextTransformer(vectorStore, SearchRequest.defaults())))
.withAugmentors(List.of(new QAPromptContextTransformer()))
.build();
Prompt prompt = new Prompt(new UserMessage("What bike is good for city commuting?"));
PromptContext promptContext = new PromptContext(prompt);
var agentResponse = chatAgent.call(promptContext);
System.out.println(agentResponse.getChatResponse().getResult().getOutput().getContent());
// RelevancyEvaluator relevancyEvaluator = new
// RelevancyEvaluator(this.chatClient);
// EvaluationRequest evaluationRequest = new EvaluationRequest(
// agentResponse.getPromptContext().getOriginalPrompt(),
// agentResponse.getPromptContext().getDataList(),
// agentResponse.getChatResponse());
//
// EvaluationResponse evaluationResponse =
// relevancyEvaluator.evaluate(evaluationRequest);
// System.out.println(evaluationResponse);
}
void loadData() {
JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description");
var textSplitter = new TokenTextSplitter();
vectorStore.accept(textSplitter.apply(jsonReader.get()));
}
@SpringBootConfiguration
static class Config {
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public ChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public EmbeddingClient embeddingClient(OpenAiApi openAiApi) {
return new OpenAiEmbeddingClient(openAiApi);
}
@Bean
public VectorStore vectorStore(EmbeddingClient embeddingClient) {
return new SimpleVectorStore(embeddingClient);
}
}
}

52
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java Normal file
View File

@@ -0,0 +1,52 @@
package org.springframework.ai.chat.agent;
import org.springframework.ai.chat.ChatResponse;
import java.util.Objects;
/**
* Encapsulates the response from the ChatAgent. Contains the most up-to-date
* PromptContext and the final ChatResponse
*
* @author Mark Pollack
* @since 1.0 M1
*/
public class AgentResponse {
private final PromptContext promptContext;
private final ChatResponse chatResponse;
public AgentResponse(PromptContext promptContext, ChatResponse chatResponse) {
this.promptContext = promptContext;
this.chatResponse = chatResponse;
}
public PromptContext getPromptContext() {
return promptContext;
}
public ChatResponse getChatResponse() {
return chatResponse;
}
@Override
public String toString() {
return "AgentResponse{" + "promptContext=" + promptContext + ", chatResponse=" + chatResponse + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof AgentResponse that))
return false;
return Objects.equals(promptContext, that.promptContext) && Objects.equals(chatResponse, that.chatResponse);
}
@Override
public int hashCode() {
return Objects.hash(promptContext, chatResponse);
}
}

21
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java Normal file
View File

@@ -0,0 +1,21 @@
package org.springframework.ai.chat.agent;
/**
* A ChatAgent encapsulates common AI workflows such as Retrieval Augmented Generation.
*
* @author Mark Pollack
* @since 1.0 M1
*/
public interface ChatAgent {
/**
* Call the chat agent to execute a workflow
* @param promptContext A shared data structure that can be used in components that
* implement the workflow. Contains the initial Prompt and a conversation ID at the
* start of the workflow.
* @return the AgentResponse that contains the ChatResponse and the latest
* PromptContext
*/
AgentResponse call(PromptContext promptContext);
}

7
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java Normal file
View File

@@ -0,0 +1,7 @@
package org.springframework.ai.chat.agent;
public interface ChatAgentListener {
void onComplete(AgentResponse agentResponse);
}

111
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java Normal file
View File

@@ -0,0 +1,111 @@
package org.springframework.ai.chat.agent;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.agent.transformer.PromptContextTransformer;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
public class DefaultChatAgent implements ChatAgent {
private ChatClient chatClient;
private List<PromptContextTransformer> retrievers;
private List<PromptContextTransformer> documentPostProcessors;
private List<PromptContextTransformer> augmentors;
private List<ChatAgentListener> chatAgentListeners;
public DefaultChatAgent(ChatClient chatClient, List<PromptContextTransformer> retrievers,
List<PromptContextTransformer> documentPostProcessors, List<PromptContextTransformer> augmentors,
List<ChatAgentListener> chatAgentListeners) {
Objects.requireNonNull(chatClient, "chatClient must not be null");
this.chatClient = chatClient;
this.retrievers = retrievers;
this.documentPostProcessors = documentPostProcessors;
this.augmentors = augmentors;
this.chatAgentListeners = chatAgentListeners;
}
public static DefaultChatAgentBuilder builder(ChatClient chatClient) {
return new DefaultChatAgentBuilder().withChatClient(chatClient);
}
@Override
public AgentResponse call(PromptContext promptContext) {
// Perform retrieval of documents and messages
for (PromptContextTransformer retriever : retrievers) {
promptContext = retriever.transform(promptContext);
}
// Perform post procesing of all retrieved documents and messages
for (PromptContextTransformer documentPostProcessor : documentPostProcessors) {
promptContext = documentPostProcessor.transform(promptContext);
}
// Perform prompt augmentation
for (PromptContextTransformer augmentor : augmentors) {
promptContext = augmentor.transform(promptContext);
}
// Perform generation
ChatResponse chatResponse = chatClient.call(promptContext.getPrompt());
// Invoke Listeners onComplete
AgentResponse agentResponse = new AgentResponse(promptContext, chatResponse);
for (ChatAgentListener listener : chatAgentListeners) {
listener.onComplete(agentResponse);
}
return agentResponse;
}
public static class DefaultChatAgentBuilder {
private ChatClient chatClient;
private List<PromptContextTransformer> retrievers = new ArrayList<>();
private List<PromptContextTransformer> documentPostProcessors = new ArrayList<>();
private List<PromptContextTransformer> augmentors = new ArrayList<>();
private List<ChatAgentListener> chatAgentListeners = new ArrayList<>();
public DefaultChatAgentBuilder withChatClient(ChatClient chatClient) {
this.chatClient = chatClient;
return this;
}
public DefaultChatAgentBuilder withRetrievers(List<PromptContextTransformer> retrievers) {
this.retrievers = retrievers;
return this;
}
public DefaultChatAgentBuilder withDocumentPostProcessors(
List<PromptContextTransformer> documentPostProcessors) {
this.documentPostProcessors = documentPostProcessors;
return this;
}
public DefaultChatAgentBuilder withAugmentors(List<PromptContextTransformer> augmentors) {
this.augmentors = augmentors;
return this;
}
public DefaultChatAgentBuilder withChatAgentListeners(List<ChatAgentListener> chatAgentListeners) {
this.chatAgentListeners = chatAgentListeners;
return this;
}
public DefaultChatAgent build() {
return new DefaultChatAgent(chatClient, retrievers, documentPostProcessors, augmentors, chatAgentListeners);
}
}
}

104
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/PromptContext.java Normal file
View File

@@ -0,0 +1,104 @@
package org.springframework.ai.chat.agent;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.node.Node;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
/**
* The shared, at the moment, mutable, data structure that can be used to implement
* ChatAgent functionality.
*
* @author Mark Pollack
* @since 1.0 M1
*/
public class PromptContext {
private Prompt prompt; // The most up-to-date prompt to use
List<Node<?>> nodes; // The most up-to-date data to use
private List<Prompt> promptHistory;
private String conversationId;
private Map<String, Object> metadata;
public PromptContext(Prompt prompt) {
this(prompt, new ArrayList<>());
}
public PromptContext(Prompt prompt, String conversationId) {
this(prompt, new ArrayList<>());
this.conversationId = conversationId;
}
public PromptContext(Prompt prompt, List<Node<?>> nodes) {
this.prompt = prompt;
this.promptHistory = new ArrayList<>();
this.promptHistory.add(prompt);
this.nodes = nodes;
}
public Prompt getPrompt() {
return prompt;
}
public void setPrompt(Prompt prompt) {
this.prompt = prompt;
}
public void addData(Node<?> datum) {
this.nodes.add(datum);
}
public List<Node<?>> getNodes() {
return nodes;
}
public void setNodes(List<Node<?>> nodes) {
this.nodes = nodes;
}
public void addPromptHistory(Prompt prompt) {
this.promptHistory.add(prompt);
}
public List<Prompt> getPromptHistory() {
return promptHistory;
}
public String getConversationId() {
return conversationId;
}
public Map<String, Object> getMetadata() {
return metadata;
}
@Override
public String toString() {
return "PromptContext{" + "prompt=" + prompt + ", nodes=" + nodes + ", promptHistory=" + promptHistory
+ ", conversationId='" + conversationId + '\'' + ", metadata=" + metadata + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof PromptContext that))
return false;
return Objects.equals(prompt, that.prompt) && Objects.equals(nodes, that.nodes)
&& Objects.equals(promptHistory, that.promptHistory)
&& Objects.equals(conversationId, that.conversationId) && Objects.equals(metadata, that.metadata);
}
@Override
public int hashCode() {
return Objects.hash(prompt, nodes, promptHistory, conversationId, metadata);
}
}

22
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/transformer/PromptContextTransformer.java Normal file
View File

@@ -0,0 +1,22 @@
package org.springframework.ai.chat.agent.transformer;
import org.springframework.ai.chat.agent.PromptContext;
/**
* Transforms the PromptContext. Implementations may retrieve data and modify the Prompt
* as needed.
*
* @author Mark Pollack
* @since 1.0 M1
*/
@FunctionalInterface
public interface PromptContextTransformer {
/**
* Transforms the given PromptContext.
* @param context the PromptContext to transform
* @return the transformed PromptContext
*/
PromptContext transform(PromptContext context);
}

74
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/transformer/QAPromptContextTransformer.java Normal file
View File

@@ -0,0 +1,74 @@
package org.springframework.ai.chat.agent.transformer;
import org.springframework.ai.chat.agent.PromptContext;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.node.Node;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
/**
* Transforms the PromptContext by adding to the prompt a Question and Answer text that
* contains the placeholder names "context" and "question".
*/
public class QAPromptContextTransformer implements PromptContextTransformer {
private static final String DEFAULT_USER_PROMPT_TEXT = """
"Context information is below.\\n"
"---------------------\\n"
"{context}\\n"
"---------------------\\n"
"Given the context information and not prior knowledge, "
"answer the question. If the answer is not in the context, inform "
"the user that you can't answer the question.\\n"
"Question: {question}\\n"
"Answer: "
""";
@Override
public PromptContext transform(PromptContext promptContext) {
String context = doCreateContext(promptContext.getNodes());
Map<String, Object> contextMap = doCreateContextMap(promptContext.getPrompt(), context);
Prompt prompt = doCreatePrompt(promptContext.getPrompt(), contextMap);
promptContext.setPrompt(prompt);
promptContext.addPromptHistory(prompt);
// For now return the modified instance instead of a copy
return promptContext;
}
protected String doCreateContext(List<Node<?>> data) {
return data.stream()
.filter(node -> node instanceof Document)
.map(node -> (Document) node)
.map(Node::getContent)
.collect(Collectors.joining(System.lineSeparator()));
}
private Map<String, Object> doCreateContextMap(Prompt prompt, String context) {
String originalUserMessage = prompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.USER)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
return Map.of("context", context, "question", originalUserMessage);
}
protected Prompt doCreatePrompt(Prompt originalPrompt, Map<String, Object> contextMap) {
PromptTemplate promptTemplate = new PromptTemplate(DEFAULT_USER_PROMPT_TEXT);
Message userMessageToAppend = promptTemplate.createMessage(contextMap);
List<Message> messageList = originalPrompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() != MessageType.USER)
.collect(Collectors.toList());
messageList.add(userMessageToAppend);
return new Prompt(messageList, (ChatOptions) originalPrompt.getOptions());
}
}

70
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/transformer/VectorStorePromptContextTransformer.java Normal file
View File

@@ -0,0 +1,70 @@
package org.springframework.ai.chat.agent.transformer;
import org.springframework.ai.chat.agent.PromptContext;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
/**
* Transforms the PromptContext by retrieving documents from a VectorStore
*/
public class VectorStorePromptContextTransformer implements PromptContextTransformer {
private final VectorStore vectorStore;
private final SearchRequest searchRequest;
public VectorStorePromptContextTransformer(VectorStore vectorStore, SearchRequest searchRequest) {
this.vectorStore = vectorStore;
this.searchRequest = searchRequest;
}
public VectorStore getVectorStore() {
return vectorStore;
}
public SearchRequest getSearchRequest() {
return searchRequest;
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Message> instructions = promptContext.getPrompt().getInstructions();
String userMessage = instructions.stream()
.filter(m -> m.getMessageType() == MessageType.USER)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
List<Document> documents = vectorStore.similaritySearch(searchRequest.withQuery(userMessage));
for (Document document : documents) {
promptContext.addData(document);
}
return promptContext;
}
@Override
public String toString() {
return "VectorStorePromptContextTransformer{" + "vectorStore=" + vectorStore + ", searchRequest="
+ searchRequest + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof VectorStorePromptContextTransformer that))
return false;
return Objects.equals(vectorStore, that.vectorStore) && Objects.equals(searchRequest, that.searchRequest);
}
@Override
public int hashCode() {
return Objects.hash(vectorStore, searchRequest);
}
}