Implement LLMResult generate(Prompt... prompts) for Azure OpenAI

This commit is contained in:
Mark Pollack
2023-08-06 20:21:18 -04:00
parent 628e697a84
commit 47cbb07a84
11 changed files with 186 additions and 73 deletions

View File

@@ -16,24 +16,20 @@
package org.springframework.ai.azure.openai.llm;
import java.util.List;
import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.models.ChatChoice;
import com.azure.ai.openai.models.ChatCompletions;
import com.azure.ai.openai.models.ChatCompletionsOptions;
import com.azure.ai.openai.models.ChatMessage;
import com.azure.ai.openai.models.ChatRole;
import com.azure.ai.openai.models.Choice;
import com.azure.ai.openai.models.Completions;
import com.azure.ai.openai.models.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.core.llm.LLMResult;
import org.springframework.ai.core.llm.LlmClient;
import org.springframework.ai.core.prompt.Generation;
import org.springframework.ai.core.prompt.Prompt;
import org.springframework.ai.core.prompt.messages.Message;
import org.springframework.util.Assert;
import java.util.ArrayList;
import java.util.List;
/**
* Implementation of {@link LlmClient} backed by an OpenAiService
*/
@@ -54,9 +50,9 @@ public class AzureOpenAiClient implements LlmClient {
@Override
public String generate(String text) {
ChatMessage chatMessage = new ChatMessage(ChatRole.USER, text);
ChatMessage azureChatMessage = new ChatMessage(ChatRole.USER, text);
ChatCompletionsOptions options = new ChatCompletionsOptions(List.of(chatMessage));
ChatCompletionsOptions options = new ChatCompletionsOptions(List.of(azureChatMessage));
options.setTemperature(this.getTemperature());
options.setModel(this.getModel());
@@ -72,7 +68,31 @@ public class AzureOpenAiClient implements LlmClient {
@Override
public LLMResult generate(Prompt... prompts) {
throw new RuntimeException("Method LLMResult generate(Prompt... prompts) not implemented.");
List<List<Generation>> generationsList = new ArrayList<>();
for (Prompt prompt : prompts) {
List<Message> messages = prompt.getMessages();
List<ChatMessage> azureMessages = new ArrayList<>();
for (Message message : messages) {
String messageType = message.getMessageType().getValue();
ChatRole chatRole = ChatRole.fromString(messageType);
azureMessages.add(new ChatMessage(chatRole, message.getContent()));
}
ChatCompletionsOptions options = new ChatCompletionsOptions(azureMessages);
options.setTemperature(this.getTemperature());
options.setModel(this.getModel());
ChatCompletions chatCompletions = this.msoftOpenAiClient.getChatCompletions(this.getModel(), options);
List<Generation> generations = new ArrayList<>();
for (ChatChoice choice : chatCompletions.getChoices()) {
ChatMessage choiceMessage = choice.getMessage();
// TODO investigate mapping of additional metadata/runtime info to the
// general model.
Generation generation = new Generation(choiceMessage.getContent());
generations.add(generation);
}
generationsList.add(generations);
}
return new LLMResult(generationsList);
}
public Double getTemperature() {

View File

@@ -33,6 +33,7 @@
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -1,41 +1,71 @@
package org.springframework.ai.core.llm;/*
* Copyright 2023 the original author or authors.
*
* 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
*
* https://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.
*/
import java.util.List;
import java.util.Map;
/*
* Copyright 2023 the original author or authors.
*
* 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
*
* https://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.
*/
package org.springframework.ai.core.llm;
import org.springframework.ai.core.prompt.Generation;
public interface LLMResult {
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class LLMResult {
private final List<List<Generation>> generations;
private Map<String, Object> providerOutput = new HashMap<>();
private Map<String, Object> runInfo = new HashMap<>();
public LLMResult(List<List<Generation>> generations) {
this.generations = generations;
}
public LLMResult(List<List<Generation>> generations, Map<String, Object> providerOutput) {
this.generations = generations;
this.providerOutput = providerOutput;
}
public LLMResult(List<List<Generation>> generations, Map<String, Object> providerOutput,
Map<String, Object> runInfo) {
this.generations = generations;
this.providerOutput = providerOutput;
this.runInfo = runInfo;
}
/**
* The list of generated outputs. It iss a list of lists because a single input could
* have multiple outputs, and multiple inputs could be passed in.
* @return
*/
List<List<Generation>> getGenerations();
public List<List<Generation>> getGenerations() {
return this.generations;
}
/**
* Arbitrary LLM-provider specific output
*/
Map<String, Object> getProviderOutput();
public Map<String, Object> getProviderOutput() {
return null;
}
/**
* The run metadata information
*/
Map<String, Object> getRunInfo();
public Map<String, Object> getRunInfo() {
return null;
}
}

View File

@@ -16,19 +16,19 @@
package org.springframework.ai.core.prompt;
import java.util.Map;
import org.springframework.ai.core.prompt.messages.AssistantMessage;
public class AiPromptTemplate extends PromptTemplate {
import java.util.Map;
public class AssistantPromptTemplate extends PromptTemplate {
private boolean example = false;
public AiPromptTemplate(String template) {
public AssistantPromptTemplate(String template) {
super(template);
}
public AiPromptTemplate(String template, boolean example) {
public AssistantPromptTemplate(String template, boolean example) {
super(template);
this.example = example;
}

View File

@@ -16,31 +16,32 @@
package org.springframework.ai.core.prompt;
import org.springframework.ai.core.prompt.messages.ChatMessage;
import org.springframework.ai.core.prompt.messages.MessageType;
import java.util.Map;
import org.springframework.ai.core.prompt.messages.ChatMessage;
/**
* A PromptTemplate that lets you specify the role as a string should the current
* implementations and their roles not suffice for your needs.
*/
public class ChatPromptTemplate extends PromptTemplate {
private String role;
private MessageType messageType;
public ChatPromptTemplate(String template) {
public ChatPromptTemplate(MessageType messageType, String template) {
super(template);
}
public ChatPromptTemplate(String template, String role) {
super(template);
this.role = role;
this.messageType = messageType;
}
@Override
public Prompt create() {
return new Prompt(new ChatMessage(render(), this.role));
return new Prompt(new ChatMessage(this.messageType, render()));
}
@Override
public Prompt create(Map<String, Object> model) {
return new Prompt(new ChatMessage(render(model), this.role));
return new Prompt(new ChatMessage(this.messageType, render(model)));
}
}

View File

@@ -16,16 +16,30 @@
package org.springframework.ai.core.prompt;
import java.util.HashMap;
import java.util.Map;
public class Generation {
private final String text;
private Map<String, Object> info;
private Map<String, Object> info = new HashMap<>();
public Generation(String text) {
this.text = text;
}
public Generation(String text, Map<String, Object> info) {
this.text = text;
this.info = info;
}
public String getText() {
return text;
}
public Map<String, Object> getInfo() {
return info;
}
}

View File

@@ -16,6 +16,11 @@
package org.springframework.ai.core.prompt;
import org.antlr.runtime.Token;
import org.antlr.runtime.TokenStream;
import org.stringtemplate.v4.ST;
import org.stringtemplate.v4.compiler.STLexer;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
@@ -24,11 +29,6 @@ import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.antlr.runtime.Token;
import org.antlr.runtime.TokenStream;
import org.stringtemplate.v4.ST;
import org.stringtemplate.v4.compiler.STLexer;
public class PromptTemplate extends AbstractPromptTemplate {
private ST st;

View File

@@ -0,0 +1,40 @@
/*
* Copyright 2023 the original author or authors.
*
* 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
*
* https://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.
*/
package org.springframework.ai.core.prompt;
import org.springframework.ai.core.prompt.messages.SystemMessage;
import org.springframework.ai.core.prompt.messages.UserMessage;
import java.util.Map;
public class UserPromptTemplate extends PromptTemplate {
public UserPromptTemplate(String template) {
super(template);
}
@Override
public Prompt create() {
return new Prompt(new SystemMessage(render()));
}
@Override
public Prompt create(Map<String, Object> model) {
return new Prompt(new UserMessage(render(model)));
}
}

View File

@@ -21,11 +21,15 @@ import java.util.Map;
public abstract class AbstractMessage implements Message {
private String content;
protected String content;
private Map<String, Object> properties = new HashMap<>();
protected Map<String, Object> properties = new HashMap<>();
private MessageType messageType;
protected MessageType messageType;
protected AbstractMessage() {
}
protected AbstractMessage(MessageType messageType, String content) {
this.messageType = messageType;

View File

@@ -20,20 +20,23 @@ import java.util.Map;
public class ChatMessage extends AbstractMessage {
private String role;
public ChatMessage(String content, String role) {
super(MessageType.SYSTEM, content);
this.role = role;
public ChatMessage(String role, String content) {
this.messageType = MessageType.valueOf(role);
this.content = content;
}
public ChatMessage(String content, String role, Map<String, Object> properties) {
super(MessageType.SYSTEM, content, properties);
this.role = role;
public ChatMessage(String role, String content, Map<String, Object> properties) {
this.messageType = MessageType.valueOf(role);
this.content = content;
this.properties = properties;
}
public String getRole() {
return role;
public ChatMessage(MessageType messageType, String content) {
super(messageType, content);
}
public ChatMessage(MessageType messageType, String content, Map<String, Object> properties) {
super(messageType, content, properties);
}
}

View File

@@ -16,13 +16,13 @@
package org.springframework.ai.core.prompt;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
@SuppressWarnings("unchecked")