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

client refactoring

Browse Source
This commit is contained in:
Mark Pollack
2023-08-15 17:46:59 -04:00
parent 1c41e648f4
commit 0e0c4b9ba7
8 changed files with 56 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/llm/AzureOpenAiClient.java
View File

@@ -20,8 +20,8 @@ import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.models.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.core.llm.LLMResponse;
import org.springframework.ai.core.llm.LlmClient;
import org.springframework.ai.core.llm.AiClient;
import org.springframework.ai.core.llm.AiResponse;
import org.springframework.ai.core.llm.Generation;
import org.springframework.ai.core.prompt.Prompt;
import org.springframework.ai.core.prompt.messages.Message;
@@ -31,9 +31,9 @@ import java.util.ArrayList;
import java.util.List;
/**
* Implementation of {@link LlmClient} backed by an OpenAiService
* Implementation of {@link AiClient} backed by an OpenAiService
*/
public class AzureOpenAiClient implements LlmClient {
public class AzureOpenAiClient implements AiClient {
private static final Logger logger = LoggerFactory.getLogger(AzureOpenAiClient.class);
@@ -67,7 +67,7 @@ public class AzureOpenAiClient implements LlmClient {
}
@Override
public LLMResponse generate(Prompt prompt) {
public AiResponse generate(Prompt prompt) {
List<Message> messages = prompt.getMessages();
List<ChatMessage> azureMessages = new ArrayList<>();
for (Message message : messages) {
@@ -87,7 +87,7 @@ public class AzureOpenAiClient implements LlmClient {
Generation generation = new Generation(choiceMessage.getContent());
generations.add(generation);
}
return new LLMResponse(generations);
return new AiResponse(generations);
}
public Double getTemperature() {

2
spring-ai-core/pom.xml
View File

@@ -29,7 +29,7 @@
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-core</artifactId>
<artifactId>spring-messaging</artifactId>
</dependency>
<dependency>

14
spring-ai-core/src/main/java/org/springframework/ai/core/llm/LlmClient.java → spring-ai-core/src/main/java/org/springframework/ai/core/llm/AiClient.java
View File

@@ -17,14 +17,16 @@
package org.springframework.ai.core.llm;
import org.springframework.ai.core.prompt.Prompt;
import org.springframework.ai.core.prompt.messages.UserMessage;
public interface LlmClient {
@FunctionalInterface
public interface AiClient {
String generate(String text);
default String generate(String message) {
Prompt prompt = new Prompt(new UserMessage(message));
return generate(prompt).getGenerations().get(0).getText();
}
// TODO Change to LLMResponse, maybe get rid of the varargs to simplify the response
// object, convenience of batch query isn't worth adding the complexity to the
// response object signature
LLMResponse generate(Prompt prompt);
AiResponse generate(Prompt prompt);
}

38
spring-ai-core/src/main/java/org/springframework/ai/core/llm/LLMResponse.java → spring-ai-core/src/main/java/org/springframework/ai/core/llm/AiResponse.java
View File

@@ -15,54 +15,58 @@
*/
package org.springframework.ai.core.llm;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class LLMResponse {
public class AiResponse {
private final List<Generation> generations;
private Map<String, Object> providerOutput = new HashMap<>();
private Map<String, Object> providerOutput;
private Map<String, Object> runInfo = new HashMap<>();
private Map<String, Object> runInfo;
public LLMResponse(List<Generation> generations) {
this.generations = generations;
public AiResponse(List<Generation> generations) {
this(generations, Collections.emptyMap(), Collections.emptyMap());
}
public LLMResponse(List<Generation> generations, Map<String, Object> providerOutput) {
this.generations = generations;
this.providerOutput = providerOutput;
public AiResponse(List<Generation> generations, Map<String, Object> providerOutput) {
this(generations, providerOutput, Collections.emptyMap());
}
public LLMResponse(List<Generation> generations, Map<String, Object> providerOutput, Map<String, Object> runInfo) {
this.generations = generations;
this.providerOutput = providerOutput;
this.runInfo = runInfo;
public AiResponse(List<Generation> generations, Map<String, Object> providerOutput, Map<String, Object> runInfo) {
this.generations = List.copyOf(generations);
this.providerOutput = Map.copyOf(providerOutput);
this.runInfo = Map.copyOf(runInfo);
}
/**
* The list of generated outputs. It is a list of lists because a single input could
* have multiple outputs.
* The list of generated outputs. It is a list of lists because the Prompt could
* request multiple output generations.
* @return
*/
public List<Generation> getGenerations() {
return this.generations;
return Collections.unmodifiableList(generations);
}
public Generation getGeneration() {
return this.generations.get(0);
}
/**
* Arbitrary LLM-provider specific output
*/
public Map<String, Object> getProviderOutput() {
return null;
return Collections.unmodifiableMap(providerOutput);
}
/**
* The run metadata information
*/
public Map<String, Object> getRunInfo() {
return null;
return Collections.unmodifiableMap(runInfo);
}
}

9
spring-ai-core/src/main/java/org/springframework/ai/core/llm/Generation.java
View File

@@ -16,6 +16,7 @@
package org.springframework.ai.core.llm;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
@@ -23,10 +24,10 @@ public class Generation {
private final String text;
private Map<String, Object> info = new HashMap<>();
private Map<String, Object> info;
public Generation(String text) {
this.text = text;
this(text, Collections.emptyMap());
}
public Generation(String text, Map<String, Object> info) {
@@ -35,11 +36,11 @@ public class Generation {
}
public String getText() {
return text;
return this.text;
}
public Map<String, Object> getInfo() {
return info;
return Collections.unmodifiableMap(this.info);
}
}

8
spring-ai-core/src/main/java/org/springframework/ai/core/prompt/messages/UserMessage.java
View File

@@ -25,12 +25,12 @@ import java.util.Map;
*/
public class UserMessage extends AbstractMessage {
public UserMessage(String content) {
super(MessageType.USER, content);
public UserMessage(String message) {
super(MessageType.USER, message);
}
public UserMessage(String content, Map<String, Object> properties) {
super(MessageType.USER, content, properties);
public UserMessage(String message, Map<String, Object> properties) {
super(MessageType.USER, message, properties);
}
}

12
spring-ai-openai/src/main/java/org/springframework/ai/openai/llm/OpenAiClient.java
View File

@@ -25,17 +25,17 @@ import com.theokanning.openai.service.OpenAiService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.core.llm.LLMResponse;
import org.springframework.ai.core.llm.LlmClient;
import org.springframework.ai.core.llm.AiClient;
import org.springframework.ai.core.llm.AiResponse;
import org.springframework.ai.core.prompt.Prompt;
import org.springframework.ai.core.prompt.messages.Message;
import org.springframework.util.Assert;
/**
* Implementation of {@link LlmClient} backed by an OpenAiService
* Implementation of {@link AiClient} backed by an OpenAiService
*/
public class OpenAiClient implements LlmClient {
public class OpenAiClient implements AiClient {
private static final Logger logger = LoggerFactory.getLogger(OpenAiClient.class);
@@ -75,7 +75,7 @@ public class OpenAiClient implements LlmClient {
}
@Override
public LLMResponse generate(Prompt prompt) {
public AiResponse generate(Prompt prompt) {
List<ChatCompletionRequest> chatCompletionRequests = getChatCompletionRequest(prompt);
return getLLMResult(chatCompletionRequests);
}
@@ -99,7 +99,7 @@ public class OpenAiClient implements LlmClient {
return response;
}
private LLMResponse getLLMResult(List<ChatCompletionRequest> chatCompletionRequest) {
private AiResponse getLLMResult(List<ChatCompletionRequest> chatCompletionRequest) {
// TODO
throw new RuntimeException("LLMResult getLLMResult not yet implemented");
}

13
spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIntegrationTest.java
View File

@@ -2,13 +2,10 @@ package org.springframework.ai.openai.acme;
import org.junit.jupiter.api.Test;
import org.springframework.ai.core.document.Document;
import org.springframework.ai.core.llm.LLMResponse;
import org.springframework.ai.core.llm.LlmClient;
import org.springframework.ai.core.llm.AiResponse;
import org.springframework.ai.core.llm.AiClient;
import org.springframework.ai.core.loader.impl.JsonLoader;
import org.springframework.ai.core.prompt.ChatPromptTemplate;
import org.springframework.ai.core.prompt.Prompt;
import org.springframework.ai.core.prompt.PromptTemplate;
import org.springframework.ai.core.prompt.messages.ChatMessage;
import org.springframework.ai.core.prompt.messages.SystemMessage;
import org.springframework.ai.core.prompt.messages.UserMessage;
import org.springframework.ai.core.retriever.impl.VectorStoreRetriever;
@@ -35,13 +32,13 @@ public class AcmeIntegrationTest {
private OpenAiEmbeddingClient embeddingClient;
@Autowired
private LlmClient llmClient;
private AiClient aiClient;
@Test
void beanTest() {
assertThat(resource).isNotNull();
assertThat(embeddingClient).isNotNull();
assertThat(llmClient).isNotNull();
assertThat(aiClient).isNotNull();
}
void acmeChain() {
@@ -74,7 +71,7 @@ public class AcmeIntegrationTest {
// Create the prompt ad-hoc for now, need to put in system message and user
// message via ChatPromptTemplate or some other message building mechanic
Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
LLMResponse response = llmClient.generate(prompt);
AiResponse response = aiClient.generate(prompt);
// Chain
// qa = new ConversationalRetrievalChain(llmClient, userPromptTemplate,