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

Replace the theokanning with custom OpenAi API client

Browse Source 
- Introducing a new OpenAiApi native client for OpenAI API and get rid of the theokanning library.
   Amongst others the OpenAiApi allows:
    - easy base-url configuration (e.g. TAS-AI)
    - Flux response for streaming OpenAI results.
    - Exposes the http headers containing important metadata
    - Pure Spring ecosystem, making it easier for  Graal VM
 - Define a new  AiStreamClient interface returning Flux<AiResponse>
 - Refactor OpenAiClient and to use the new OpenAiApi and implement the AiStreamClient.
 - Use spring-retry to improve the OpenAI EmbeddingClient stability on 503 error.

 - Remove the OpenAI http header interceptor as the OpenAiApi returns ResponseEntity<T> that provides direct access to the headers.
 - Refactor the metadata headers and usage extraction.
 - Remove redundant and obsolete classes.
 - Fix dependency issue with Pinecone, netty-codec-http2 and Spring Boot 3.2
This commit is contained in:
Christian Tzolov
2023-12-13 15:48:28 +01:00
parent 7aad51da7f
commit e30be94a56
37 changed files with 1449 additions and 1190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pom.xml
View File

@@ -84,8 +84,8 @@
<!-- production dependencies -->
<spring-boot.version>3.2.0</spring-boot.version>
<spring-framework.version>6.1.1</spring-framework.version>
<stringtemplate.version>4.0.2</stringtemplate.version>
<open-ai-client.version>0.16.0</open-ai-client.version>
<azure-open-ai-client.version>1.0.0-beta.3</azure-open-ai-client.version>
<jtokkit.version>0.6.1</jtokkit.version>
<victools.version>4.31.1</victools.version>

3
spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/client/AzureOpenAiClientMetadataTests.java
View File

@@ -17,7 +17,6 @@
package org.springframework.ai.azure.openai.client;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.test.config.MockAiTestConfiguration.SPRING_AI_API_PATH;
import java.nio.charset.StandardCharsets;
@@ -170,7 +169,7 @@ class AzureOpenAiClientMetadataTests {
}
@RestController
@RequestMapping(SPRING_AI_API_PATH)
@RequestMapping("/spring-ai/api")
@SuppressWarnings("all")
static class SpringAzureOpenAiChatCompletionsController {

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

@@ -42,7 +42,7 @@
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>6.1.1</version>
<version>${spring-framework.version}</version>
</dependency>
<dependency>

11
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/AiStreamClient.java → spring-ai-core/src/main/java/org/springframework/ai/client/AiStreamClient.java
View File

@@ -14,20 +14,15 @@
* limitations under the License.
*/
package org.springframework.ai.openai.client;
package org.springframework.ai.client;
import reactor.core.publisher.Flux;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.messages.UserMessage;
@FunctionalInterface
public interface AiStreamClient {
default Flux<OpenAiSseResponse> generateStream(String message) {
Prompt prompt = new Prompt(new UserMessage(message));
return generateStream(prompt);
}
Flux<OpenAiSseResponse> generateStream(Prompt prompt);
public Flux<AiResponse> generateStream(Prompt prompt);
}

12
spring-ai-openai/pom.xml
View File

@@ -27,6 +27,12 @@
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.retry</groupId>
<artifactId>spring-retry</artifactId>
<version>2.0.4</version>
</dependency>
<dependency>
<groupId>io.rest-assured</groupId>
<artifactId>json-path</artifactId>
@@ -37,12 +43,6 @@
<artifactId>okhttp</artifactId>
</dependency>
<dependency>
<groupId>com.theokanning.openai-gpt3-java</groupId>
<artifactId>service</artifactId>
<version>${open-ai-client.version}</version>
</dependency>
<dependency>
<groupId>com.github.victools</groupId>
<artifactId>jsonschema-generator</artifactId>

718
spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java Normal file
View File

@@ -0,0 +1,718 @@
/*
* Copyright 2023-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.openai.api;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
// @formatter:off
/**
* Single class implementation of the OpenAI Chat Completion API: https://beta.openai.com/docs/api-reference/chat and
* OpenAI Embedding API: https://beta.openai.com/docs/api-reference/embeddings.
*
* @author Christian Tzolov
*/
public class OpenAiApi {
private static final String DEFAULT_BASE_URL = "https://api.openai.com";
private static final String DEFAULT_EMBEDDING_MODEL = "text-embedding-ada-002";
private static final String SSE_DONE = "[DONE]";
private final RestClient restClient;
private final WebClient webClient;
private final ObjectMapper objectMapper;
/**
* Create an new chat completion api with base URL set to https://api.openai.com
*
* @param openAiToken OpenAI apiKey.
*/
public OpenAiApi(String openAiToken) {
this(DEFAULT_BASE_URL, openAiToken, RestClient.builder());
}
/**
* Create an new chat completion api.
*
* @param baseUrl api base URL.
* @param openAiToken OpenAI apiKey.
* @param restClientBuilder RestClient builder.
*/
public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder) {
this.objectMapper = new ObjectMapper();
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
headers.setBearerAuth(openAiToken);
headers.setContentType(MediaType.APPLICATION_JSON);
};
var responseErrorHandler = new ResponseErrorHandler() {
@Override
public boolean hasError(ClientHttpResponse response) throws IOException {
return response.getStatusCode().isError();
}
@Override
public void handleError(ClientHttpResponse response) throws IOException {
if (response.getStatusCode().isError()) {
throw new OpenAiApiException(String.format("%s - %s", response.getStatusCode().value(),
new ObjectMapper().readValue(response.getBody(), ResponseError.class)));
}
}
};
this.restClient = restClientBuilder
.baseUrl(baseUrl)
.defaultHeaders(jsonContentHeaders)
.defaultStatusHandler(responseErrorHandler)
.build();
this.webClient = WebClient.builder()
.baseUrl(baseUrl)
.defaultHeaders(jsonContentHeaders)
.build();
}
public static class OpenAiApiException extends RuntimeException {
public OpenAiApiException(String message) {
super(message);
}
public OpenAiApiException(String message, Throwable cause) {
super(message, cause);
}
}
/**
* API error response.
* @param error Error details.
*/
@JsonInclude(Include.NON_NULL)
public record ResponseError(@JsonProperty("error") Error error) {
/**
* Error details.
* @param message Error message.
* @param type Error type.
* @param param Error parameter.
* @param code Error code.
*/
@JsonInclude(Include.NON_NULL)
public record Error(
@JsonProperty("message") String message,
@JsonProperty("type") String type,
@JsonProperty("param") String param,
@JsonProperty("code") String code) {
}
}
/**
* Represents a tool the model may call. Currently, only functions are supported as a tool.
*
* @param type The type of the tool. Currently, only 'function' is supported.
* @param function The function definition.
*/
@JsonInclude(Include.NON_NULL)
public record FunctionTool(
@JsonProperty("type") Type type,
@JsonProperty("function") Function function) {
/**
* Create a tool of type 'function' and the given function definition.
* @param function function definition.
*/
public FunctionTool(Function function) {
this(Type.function, function);
}
/**
* Create a tool of type 'function' and the given function definition.
*/
public enum Type {
/**
* Function tool type.
*/
function
}
/**
* Function definition.
*
* @param description A description of what the function does, used by the model to choose when and how to call
* the function.
* @param name The name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes,
* with a maximum length of 64.
* @param parameters The parameters the functions accepts, described as a JSON Schema object. To describe a
* function that accepts no parameters, provide the value {"type": "object", "properties": {}}.
*/
public record Function(
@JsonProperty("description") String description,
@JsonProperty("name") String name,
@JsonProperty("parameters") Map<String, Object> parameters) {
/**
* Create tool function definition.
*
* @param description tool function description.
* @param name tool function name.
* @param jsonSchema tool function schema as json.
*/
public Function(String description, String name, String jsonSchema) {
this(description, name, parseJson(jsonSchema));
}
}
}
/**
* Creates a model response for the given chat conversation.
*
* @param messages A list of messages comprising the conversation so far.
* @param model ID of the model to use.
* @param frequencyPenalty Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing
* frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.
* @param logitBias Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object
* that maps tokens (specified by their token ID in the tokenizer) to an associated bias value from -100 to 100.
* Mathematically, the bias is added to the logits generated by the model prior to sampling. The exact effect will
* vary per model, but values between -1 and 1 should decrease or increase likelihood of selection; values like -100
* or 100 should result in a ban or exclusive selection of the relevant token.
* @param maxTokens The maximum number of tokens to generate in the chat completion. The total length of input
* tokens and generated tokens is limited by the model's context length.
* @param n How many chat completion choices to generate for each input message. Note that you will be charged based
* on the number of generated tokens across all of the choices. Keep n as 1 to minimize costs.
* @param presencePenalty Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they
* appear in the text so far, increasing the model's likelihood to talk about new topics.
* @param responseFormat An object specifying the format that the model must output. Setting to { "type":
* "json_object" } enables JSON mode, which guarantees the message the model generates is valid JSON.
* @param seed This feature is in Beta. If specified, our system will make a best effort to sample
* deterministically, such that repeated requests with the same seed and parameters should return the same result.
* Determinism is not guaranteed, and you should refer to the system_fingerprint response parameter to monitor
* changes in the backend.
* @param stop Up to 4 sequences where the API will stop generating further tokens.
* @param stream If set, partial message deltas will be sent.Tokens will be sent as data-only server-sent events as
* they become available, with the stream terminated by a data: [DONE] message.
* @param temperature What sampling temperature to use, between 0 and 1. Higher values like 0.8 will make the output
* more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend
* altering this or top_p but not both.
* @param topP An alternative to sampling with temperature, called nucleus sampling, where the model considers the
* results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10%
* probability mass are considered. We generally recommend altering this or temperature but not both.
* @param tools A list of tools the model may call. Currently, only functions are supported as a tool. Use this to
* provide a list of functions the model may generate JSON inputs for.
* @param toolChoice Controls which (if any) function is called by the model. none means the model will not call a
* function and instead generates a message. auto means the model can pick between generating a message or calling a
* function. Specifying a particular function via {"type: "function", "function": {"name": "my_function"}} forces
* the model to call that function. none is the default when no functions are present. auto is the default if
* functions are present.
* @param user A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
*
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionRequest(
@JsonProperty("messages") List<ChatCompletionMessage> messages,
@JsonProperty("model") String model,
@JsonProperty("frequency_penalty") Float frequencyPenalty,
@JsonProperty("logit_bias") Map<String, Object> logitBias,
@JsonProperty("max_tokens") Integer maxTokens,
@JsonProperty("n") Integer n,
@JsonProperty("presence_penalty") Float presencePenalty,
@JsonProperty("response_format") ResponseFormat responseFormat,
@JsonProperty("seed") Integer seed,
@JsonProperty("stop") String stop,
@JsonProperty("stream") Boolean stream,
@JsonProperty("temperature") Float temperature,
@JsonProperty("top_p") Float topP,
@JsonProperty("tools") List<FunctionTool> tools,
@JsonProperty("tool_choice") ToolChoice toolChoice,
@JsonProperty("user") String user) {
/**
* Shortcut constructor for a chat completion request with the given messages and model.
*
* @param messages A list of messages comprising the conversation so far.
* @param model ID of the model to use.
* @param temperature What sampling temperature to use, between 0 and 1.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature) {
this(messages, model, 0.0f, null, null, 1, 0.0f,
null, null, null, false, temperature, null,
null, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages, model and control for streaming.
*
* @param messages A list of messages comprising the conversation so far.
* @param model ID of the model to use.
* @param temperature What sampling temperature to use, between 0 and 1.
* @param stream If set, partial message deltas will be sent.Tokens will be sent as data-only server-sent events
* as they become available, with the stream terminated by a data: [DONE] message.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature, boolean stream) {
this(messages, model, 0.0f, null, null, 1, 0.0f,
null, null, null, stream, temperature, null,
null, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages, model, tools and tool choice.
* Streaming is set to false, temperature to 0.8 and all other parameters are null.
*
* @param messages A list of messages comprising the conversation so far.
* @param model ID of the model to use.
* @param tools A list of tools the model may call. Currently, only functions are supported as a tool.
* @param toolChoice Controls which (if any) function is called by the model.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model,
List<FunctionTool> tools, ToolChoice toolChoice) {
this(messages, model, 0.0f, null, null, 1, 0.0f,
null, null, null, false, 0.8f, null,
tools, toolChoice, null);
}
/**
* Specifies a tool the model should use. Use to force the model to call a specific function.
*
* @param type The type of the tool. Currently, only 'function' is supported.
* @param function single field map for type 'name':'your function name'.
*/
@JsonInclude(Include.NON_NULL)
public record ToolChoice(
@JsonProperty("type") String type,
@JsonProperty("function") Map<String, String> function) {
/**
* Create a tool choice of type 'function' and name 'functionName'.
* @param functionName Function name of the tool.
*/
public ToolChoice(String functionName) {
this("function", Map.of("name", functionName));
}
}
/**
* An object specifying the format that the model must output.
* @param type Must be one of 'text' or 'json_object'.
*/
@JsonInclude(Include.NON_NULL)
public record ResponseFormat(
@JsonProperty("type") String type) {
}
}
/**
* Message comprising the conversation.
*
* @param content The contents of the message.
* @param role The role of the messages author. Could be one of the {@link Role} types.
* @param name An optional name for the participant. Provides the model information to differentiate between
* participants of the same role.
* @param toolCallId Tool call that this message is responding to. Only applicable for the {@link Role#tool} role
* and null otherwise.
* @param toolCalls The tool calls generated by the model, such as function calls. Applicable only for
* {@link Role#assistant} role and null otherwise.
* @param functionCall Deprecated and replaced by tool_calls. The name and arguments of a function that should be
* called, as generated by the model.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionMessage(
@JsonProperty("content") String content,
@JsonProperty("role") Role role,
@JsonProperty("name") String name,
@JsonProperty("tool_call_id") String toolCallId,
@JsonProperty("tool_calls") List<ToolCall> toolCalls,
@JsonProperty("function_call") ChatCompletionFunction functionCall) {
/**
* Create a chat completion message with the given content and role. All other fields are null.
* @param content The contents of the message.
* @param role The role of the author of this message.
*/
public ChatCompletionMessage(String content, Role role) {
this(content, role, null, null, null, null);
}
/**
* The role of the author of this message.
*/
public enum Role {
/**
* System message.
*/
system,
/**
* User message.
*/
user,
/**
* Assistant message.
*/
assistant,
/**
* Tool message.
*/
tool
}
/**
* The relevant tool call.
*
* @param id The ID of the tool call. This ID must be referenced when you submit the tool outputs in using the
* Submit tool outputs to run endpoint.
* @param type The type of tool call the output is required for. For now, this is always function.
* @param function The function definition.
*/
@JsonInclude(Include.NON_NULL)
public record ToolCall(
@JsonProperty("id") String id,
@JsonProperty("type") String type,
@JsonProperty("function") ChatCompletionFunction function) {
}
/**
* The function definition.
*
* @param name The name of the function.
* @param arguments The arguments that the model expects you to pass to the function.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionFunction(
@JsonProperty("name") String name,
@JsonProperty("arguments") String arguments) {
}
}
/**
* The reason the model stopped generating tokens.
*/
public enum ChatCompletionFinishReason {
/**
* The model hit a natural stop point or a provided stop sequence.
*/
stop,
/**
* The maximum number of tokens specified in the request was reached.
*/
length,
/**
* The content was omitted due to a flag from our content filters.
*/
content_filter,
/**
* The model called a tool.
*/
tool_calls,
/**
* (deprecated) The model called a function.
*/
function_call
}
/**
* Represents a chat completion response returned by model, based on the provided input.
*
* @param id A unique identifier for the chat completion.
* @param choices A list of chat completion choices. Can be more than one if n is greater than 1.
* @param created The Unix timestamp (in seconds) of when the chat completion was created.
* @param model The model used for the chat completion.
* @param systemFingerprint This fingerprint represents the backend configuration that the model runs with. Can be
* used in conjunction with the seed request parameter to understand when backend changes have been made that might
* impact determinism.
* @param object The object type, which is always chat.completion.
* @param usage Usage statistics for the completion request.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletion(
@JsonProperty("id") String id,
@JsonProperty("choices") List<Choice> choices,
@JsonProperty("created") Long created,
@JsonProperty("model") String model,
@JsonProperty("system_fingerprint") String systemFingerprint,
@JsonProperty("object") String object,
@JsonProperty("usage") Usage usage) {
/**
* Chat completion choice.
*
* @param finishReason The reason the model stopped generating tokens.
* @param index The index of the choice in the list of choices.
* @param message A chat completion message generated by the model.
*/
@JsonInclude(Include.NON_NULL)
public record Choice(
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason,
@JsonProperty("index") Integer index,
@JsonProperty("message") ChatCompletionMessage message) {
}
}
/**
* Usage statistics for the completion request.
*
* @param completionTokens Number of tokens in the generated completion. Only applicable for completion requests.
* @param promptTokens Number of tokens in the prompt.
* @param totalTokens Total number of tokens used in the request (prompt + completion).
*/
@JsonInclude(Include.NON_NULL)
public record Usage(
@JsonProperty("completion_tokens") Integer completionTokens,
@JsonProperty("prompt_tokens") Integer promptTokens,
@JsonProperty("total_tokens") Integer totalTokens) {
}
/**
* Represents a streamed chunk of a chat completion response returned by model, based on the provided input.
*
* @param id A unique identifier for the chat completion. Each chunk has the same ID.
* @param choices A list of chat completion choices. Can be more than one if n is greater than 1.
* @param created The Unix timestamp (in seconds) of when the chat completion was created. Each chunk has the same
* timestamp.
* @param model The model used for the chat completion.
* @param systemFingerprint This fingerprint represents the backend configuration that the model runs with. Can be
* used in conjunction with the seed request parameter to understand when backend changes have been made that might
* impact determinism.
* @param object The object type, which is always 'chat.completion.chunk'.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionChunk(
@JsonProperty("id") String id,
@JsonProperty("choices") List<ChunkChoice> choices,
@JsonProperty("created") Long created,
@JsonProperty("model") String model,
@JsonProperty("system_fingerprint") String systemFingerprint,
@JsonProperty("object") String object) {
/**
* Chat completion choice.
*
* @param finishReason The reason the model stopped generating tokens.
* @param index The index of the choice in the list of choices.
* @param delta A chat completion delta generated by streamed model responses.
*/
@JsonInclude(Include.NON_NULL)
public record ChunkChoice(
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason,
@JsonProperty("index") Integer index,
@JsonProperty("delta") ChatCompletionMessage delta) {
}
}
/**
* Creates a model response for the given chat conversation.
*
* @param chatRequest The chat completion request.
* @return Entity response with {@link ChatCompletion} as a body and HTTP status code and headers.
*/
public ResponseEntity<ChatCompletion> chatCompletionEntity(ChatCompletionRequest chatRequest) {
Assert.notNull(chatRequest, "The request body can not be null.");
Assert.isTrue(!chatRequest.stream(), "Request must set the steam property to false.");
return this.restClient.post()
.uri("/v1/chat/completions")
.body(chatRequest)
.retrieve()
.toEntity(ChatCompletion.class);
}
/**
* Creates a streaming chat response for the given chat conversation.
*
* @param chatRequest The chat completion request. Must have the stream property set to true.
* @return Returns a {@link Flux} stream from chat completion chunks.
*/
public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest) {
Assert.notNull(chatRequest, "The request body can not be null.");
Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true.");
return this.webClient.post()
.uri("/v1/chat/completions")
.body(Mono.just(chatRequest), ChatCompletionRequest.class)
.retrieve()
.bodyToFlux(String.class)
// cancels the flux stream after the SSE_DONE is received.
.takeUntil(content -> content.contains(SSE_DONE))
// filters out the SSE_DONE message.
.filter(content -> !content.contains(SSE_DONE))
.map(content -> parseJson(content, ChatCompletionChunk.class));
}
/**
* Represents an embedding vector returned by embedding endpoint.
*
* @param index The index of the embedding in the list of embeddings.
* @param embedding The embedding vector, which is a list of floats. The length of vector depends on the model.
* @param object The object type, which is always 'embedding'.
*/
@JsonInclude(Include.NON_NULL)
public record Embedding(
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> embedding,
@JsonProperty("object") String object) {
/**
* Create an embedding with the given index, embedding and object type set to 'embedding'.
*
* @param index The index of the embedding in the list of embeddings.
* @param embedding The embedding vector, which is a list of floats. The length of vector depends on the model.
*/
public Embedding(Integer index, List<Double> embedding) {
this(index, embedding, "embedding");
}
}
/**
* Creates an embedding vector representing the input text.
*
* @param input Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single
* request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for
* the model (8192 tokens for text-embedding-ada-002), cannot be an empty string, and any array must be 2048
* dimensions or less.
* @param model ID of the model to use.
* @param encodingFormat The format to return the embeddings in. Can be either float or base64.
* @param user A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingRequest<T>(
@JsonProperty("input") T input,
@JsonProperty("model") String model,
@JsonProperty("encoding_format") String encodingFormat,
@JsonProperty("user") String user) {
/**
* Create an embedding request with the given input, model and encoding format set to float.
* @param input Input text to embed.
* @param model ID of the model to use.
*/
public EmbeddingRequest(T input, String model) {
this(input, model, "float", null);
}
/**
* Create an embedding request with the given input. Encoding format is set to float and user is null and the
* model is set to 'text-embedding-ada-002'.
* @param input Input text to embed.
*/
public EmbeddingRequest(T input) {
this(input, DEFAULT_EMBEDDING_MODEL);
}
}
/**
* List of multiple embedding responses.
*
* @param <T> Type of the entities in the data list.
* @param object Must have value "list".
* @param data List of entities.
* @param model ID of the model to use.
* @param usage Usage statistics for the completion request.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingList<T>(
@JsonProperty("object") String object,
@JsonProperty("data") List<T> data,
@JsonProperty("model") String model,
@JsonProperty("usage") Usage usage) {
}
/**
* Creates an embedding vector representing the input text or token array.
*
* @param embeddingRequest The embedding request.
* @return Returns list of {@link Embedding} wrapped in {@link EmbeddingList}.
* @param <T> Type of the entity in the data list. Can be a {@link String} or {@link List} of tokens (e.g.
* Integers). For embedding multiple inputs in a single request, You can pass a {@link List} of {@link String} or
* {@link List} of {@link List} of tokens. For example:
*
* <pre>{@code List.of("text1", "text2", "text3") or List.of(List.of(1, 2, 3), List.of(3, 4, 5))} </pre>
*/
public <T> ResponseEntity<EmbeddingList<Embedding>> embeddings(EmbeddingRequest<T> embeddingRequest) {
Assert.notNull(embeddingRequest, "The request body can not be null.");
// Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single
// request, pass an array of strings or array of token arrays.
Assert.notNull(embeddingRequest.input(), "The input can not be null.");
Assert.isTrue(embeddingRequest.input() instanceof String || embeddingRequest.input() instanceof List,
"The input must be either a String, or a List of Strings or List of List of integers.");
// The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002), cannot
// be an empty string, and any array must be 2048 dimensions or less.
if (embeddingRequest.input() instanceof List list) {
Assert.isTrue(!CollectionUtils.isEmpty(list), "The input list can not be empty.");
Assert.isTrue(list.size() <= 2048, "The list must be 2048 dimensions or less");
Assert.isTrue(list.get(0) instanceof String || list.get(0) instanceof Integer
|| list.get(0) instanceof List,
"The input must be either a String, or a List of Strings or list of list of integers.");
}
return this.restClient.post()
.uri("/v1/embeddings")
.body(embeddingRequest)
.retrieve()
.toEntity(new ParameterizedTypeReference<>() {
});
}
private static Map<String, Object> parseJson(String jsonSchema) {
try {
return new ObjectMapper().readValue(jsonSchema,
new TypeReference<Map<String, Object>>() {
});
}
catch (Exception e) {
throw new OpenAiApiException("Failed to parse schema: " + jsonSchema, e);
}
}
private <T> T parseJson(String json, Class<T> type) {
try {
return this.objectMapper.readValue(json, type);
}
catch (Exception e) {
throw new OpenAiApiException("Failed to parse schema: " + json, e);
}
}
}
// @formatter:on

31
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionResponse.java
View File

@@ -1,31 +0,0 @@
package org.springframework.ai.openai.client;
import com.fasterxml.jackson.annotation.JsonProperty;
import java.util.List;
import java.util.Map;
public record ChatCompletionResponse(
@JsonProperty("id") String id,
@JsonProperty("choices") List<Choice> choices,
@JsonProperty("created") Integer created,
@JsonProperty("model") String model,
@JsonProperty("system_fingerprint") String systemFingerprint,
@JsonProperty("object") String object,
@JsonProperty("usage") Map<String, Object> usage) {
public record Choice(
@JsonProperty("finish_reason") String finishReason,
@JsonProperty("index") Integer index,
@JsonProperty("message") OpenAiChatMessage message) {
}
}

286
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionsRequest.java
View File

@@ -1,286 +0,0 @@
package org.springframework.ai.openai.client;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import java.util.List;
import java.util.Map;
@JsonDeserialize(builder = ChatCompletionsRequest.Builder.class)
public class ChatCompletionsRequest {
private final List<OpenAiChatMessage> messages;
private final String model;
private final Integer frequencyPenalty;
private final Map<String, String> logitBias;
private final Integer maxTokens;
private final Integer n;
private final Integer presencePenalty;
private final ResponseFormat responseFormat;
private final Integer seed;
private final List<String> stop;
private final Boolean stream;
private final Double temperature;
private final Integer topP;
private final List<Tool> tools;
private final String toolChoice;
private final String user;
private ChatCompletionsRequest(Builder builder) {
this.messages = builder.messages;
this.model = builder.model;
this.frequencyPenalty = builder.frequencyPenalty;
this.logitBias = builder.logitBias;
this.maxTokens = builder.maxTokens;
this.n = builder.n;
this.presencePenalty = builder.presencePenalty;
this.responseFormat = builder.responseFormat;
this.seed = builder.seed;
this.stop = builder.stop;
this.stream = builder.stream;
this.temperature = builder.temperature;
this.topP = builder.topP;
this.tools = builder.tools;
this.toolChoice = builder.toolChoice;
this.user = builder.user;
}
public List<OpenAiChatMessage> getMessages() {
return messages;
}
public String getModel() {
return model;
}
public Integer getFrequencyPenalty() {
return frequencyPenalty;
}
public Map<String, String> getLogitBias() {
return logitBias;
}
public Integer getMaxTokens() {
return maxTokens;
}
public Integer getN() {
return n;
}
public Integer getPresencePenalty() {
return presencePenalty;
}
public ResponseFormat getResponseFormat() {
return responseFormat;
}
public Integer getSeed() {
return seed;
}
public List<String> getStop() {
return stop;
}
public Boolean getStream() {
return stream;
}
public Double getTemperature() {
return temperature;
}
public Integer getTopP() {
return topP;
}
public List<Tool> getTools() {
return tools;
}
public String getToolChoice() {
return toolChoice;
}
public String getUser() {
return user;
}
public static class Builder {
@JsonProperty("messages")
private List<OpenAiChatMessage> messages;
@JsonProperty("model")
private String model;
@JsonProperty("frequency_penalty")
private Integer frequencyPenalty;
@JsonProperty("logit_bias")
private Map<String, String> logitBias;
@JsonProperty("max_tokens")
private Integer maxTokens;
@JsonProperty("n")
private Integer n;
@JsonProperty("presence_penalty")
private Integer presencePenalty;
@JsonProperty("response_format")
private ResponseFormat responseFormat;
@JsonProperty("seed")
private Integer seed;
@JsonProperty("stop")
private List<String> stop;
@JsonProperty("stream")
private Boolean stream;
@JsonProperty("temperature")
private Double temperature;
@JsonProperty("top_p")
private Integer topP;
@JsonProperty("tools")
private List<Tool> tools;
@JsonProperty("tool_choice")
private String toolChoice;
@JsonProperty("user")
private String user;
public Builder messages(List<OpenAiChatMessage> messages) {
this.messages = messages;
return this;
}
public Builder model(String model) {
this.model = model;
return this;
}
public Builder frequencyPenalty(Integer frequencyPenalty) {
this.frequencyPenalty = frequencyPenalty;
return this;
}
public Builder logitBias(Map<String, String> logitBias) {
this.logitBias = logitBias;
return this;
}
public Builder maxTokens(Integer maxTokens) {
this.maxTokens = maxTokens;
return this;
}
public Builder n(Integer n) {
this.n = n;
return this;
}
public Builder presencePenalty(Integer presencePenalty) {
this.presencePenalty = presencePenalty;
return this;
}
public Builder responseFormat(ResponseFormat responseFormat) {
this.responseFormat = responseFormat;
return this;
}
public Builder seed(Integer seed) {
this.seed = seed;
return this;
}
public Builder stop(List<String> stop) {
this.stop = stop;
return this;
}
public Builder stream(Boolean stream) {
this.stream = stream;
return this;
}
public Builder temperature(Double temperature) {
this.temperature = temperature;
return this;
}
public Builder topP(Integer topP) {
this.topP = topP;
return this;
}
public Builder tools(List<Tool> tools) {
this.tools = tools;
return this;
}
public Builder toolChoice(String toolChoice) {
this.toolChoice = toolChoice;
return this;
}
public Builder user(String user) {
this.user = user;
return this;
}
public ChatCompletionsRequest build() {
return new ChatCompletionsRequest(this);
}
}
public record Function(
@JsonProperty("name") String name,
@JsonProperty("description") String description,
@JsonProperty("parameters") Map<String, Object> parameters,
@JsonProperty("arguments") String arguments) {
}
public record ResponseFormat(
@JsonProperty("type") String type) {
}
public record Tool(
@JsonProperty("function") Function function,
@JsonProperty("type") String type) {
}
}

91
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiChatMessage.java
View File

@@ -1,91 +0,0 @@
package org.springframework.ai.openai.client;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import java.util.List;
@JsonDeserialize(builder = OpenAiChatMessage.Builder.class)
public class OpenAiChatMessage {
private final String role;
private final String name;
private final List<ToolCall> toolCalls;
private final String content;
private OpenAiChatMessage(Builder builder) {
this.role = builder.role;
this.name = builder.name;
this.toolCalls = builder.toolCalls;
this.content = builder.content;
}
public String getRole() {
return role;
}
public String getName() {
return name;
}
public List<ToolCall> getToolCalls() {
return toolCalls;
}
public String getContent() {
return content;
}
public static class Builder {
@JsonProperty("role")
private String role;
@JsonProperty("name")
private String name;
@JsonProperty("tool_calls")
private List<ToolCall> toolCalls;
@JsonProperty("content")
private String content;
public Builder role(String role) {
this.role = role;
return this;
}
public Builder name(String name) {
this.name = name;
return this;
}
public Builder toolCalls(List<ToolCall> toolCalls) {
this.toolCalls = toolCalls;
return this;
}
public Builder content(String content) {
this.content = content;
return this;
}
public OpenAiChatMessage build() {
return new OpenAiChatMessage(this);
}
}
public record ToolCall(
@JsonProperty("function") ChatCompletionsRequest.Function function,
@JsonProperty("id") String id,
@JsonProperty("type") String type) {
}
}

232
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiClient.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023 the original author or authors.
* Copyright 2023-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.
@@ -13,71 +13,58 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.client;
import com.theokanning.openai.client.OpenAiApi;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.client.Generation;
import org.springframework.ai.metadata.ChoiceMetadata;
import org.springframework.ai.openai.metadata.OpenAiGenerationMetadata;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.ai.prompt.messages.MessageType;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import org.springframework.context.annotation.ImportRuntimeHints;
import org.springframework.util.Assert;
import com.theokanning.openai.completion.chat.ChatCompletionChoice;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatCompletionResult;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.service.OpenAiService;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.client.AiStreamClient;
import org.springframework.ai.client.Generation;
import org.springframework.ai.metadata.ChoiceMetadata;
import org.springframework.ai.metadata.RateLimit;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage;
import org.springframework.ai.openai.metadata.OpenAiGenerationMetadata;
import org.springframework.ai.openai.metadata.support.OpenAiResponseHeaderExtractor;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
/**
* {@link AiClient} implementation for {@literal OpenAI} backed by {@link OpenAiService}.
* {@link AiClient} implementation for {@literal OpenAI} backed by {@link OpenAiApi}.
*
* @author Mark Pollack
* @author Christian Tzolov
* @author Ueibin Kim
* @author John Blum
* @author Josh Long
* @author Jemin Huh
* @see org.springframework.ai.client.AiClient
* @see com.theokanning.openai.service.OpenAiService
* @see org.springframework.ai.client.AiStreamClient
* @see OpenAiApi
*/
@ImportRuntimeHints(OpenAiClient.Hints.class)
public class OpenAiClient implements AiClient {
public class OpenAiClient implements AiClient, AiStreamClient {
static class Hints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
hints.proxies().registerJdkProxy(OpenAiApi.class);
}
}
// TODO how to set default options for the entire client
// TODO expose request options into Prompt API via PromptOptions
private Double temperature = 0.7;
private String model = "gpt-3.5-turbo";
private final Logger logger = LoggerFactory.getLogger(getClass());
private final OpenAiService openAiService;
private final OpenAiApi openAiApi;
public OpenAiClient(OpenAiService openAiService) {
Assert.notNull(openAiService, "OpenAiService must not be null");
this.openAiService = openAiService;
public OpenAiClient(OpenAiApi openAiApi) {
Assert.notNull(openAiApi, "OpenAiApi must not be null");
this.openAiApi = openAiApi;
}
public String getModel() {
@@ -96,109 +83,84 @@ public class OpenAiClient implements AiClient {
this.temperature = temperature;
}
@Override
public String generate(String text) {
ChatCompletionRequest chatCompletionRequest = getChatCompletionRequest(text);
return getResponse(chatCompletionRequest);
}
@Override
public AiResponse generate(Prompt prompt) {
List<Message> messages = prompt.getMessages();
List<ChatMessage> theoMessages = messages.stream()
.map(message -> new ChatMessage(message.getMessageTypeValue(), message.getContent()))
List<ChatCompletionMessage> chatCompletionMessages = messages.stream()
.map(m -> new ChatCompletionMessage(m.getContent(),
ChatCompletionMessage.Role.valueOf(m.getMessageType().getValue())))
.toList();
ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
.model(this.model)
.temperature(this.temperature)
.messages(theoMessages)
.build();
ResponseEntity<ChatCompletion> completionEntity = this.openAiApi.chatCompletionEntity(
new OpenAiApi.ChatCompletionRequest(chatCompletionMessages, this.model, this.temperature.floatValue()));
return getAiResponse(chatCompletionRequest);
}
private AiResponse getAiResponse(ChatCompletionRequest chatCompletionRequest) {
logger.trace("ChatMessages: {}", chatCompletionRequest.getMessages());
ChatCompletionResult chatCompletionResult = this.openAiService.createChatCompletion(chatCompletionRequest);
List<ChatCompletionChoice> chatCompletionChoices = chatCompletionResult.getChoices();
logger.trace("ChatCompletionChoices: {}", chatCompletionChoices);
List<Generation> generations = new ArrayList<>();
for (ChatCompletionChoice chatCompletionChoice : chatCompletionChoices) {
ChatMessage chatMessage = chatCompletionChoice.getMessage();
Generation generation = new Generation(chatMessage.getContent(), Map.of("role", chatMessage.getRole()))
.withChoiceMetadata(ChoiceMetadata.from(chatCompletionChoice.getFinishReason(), null));
generations.add(generation);
}
return new AiResponse(generations, OpenAiGenerationMetadata.from(chatCompletionResult));
}
private ChatCompletionRequest getChatCompletionRequest(String text) {
List<ChatMessage> chatMessages = List.of(new ChatMessage("user", text));
logger.trace("ChatMessages: {}", chatMessages);
ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
.model(this.model)
.temperature(this.temperature)
.messages(List.of(new ChatMessage("user", text)))
.build();
logger.trace("ChatCompletionRequest: {}", chatCompletionRequest);
return chatCompletionRequest;
}
private String getResponse(ChatCompletionRequest chatCompletionRequest) {
StringBuilder builder = new StringBuilder();
this.openAiService.createChatCompletion(chatCompletionRequest)
.getChoices()
.forEach(choice -> builder.append(choice.getMessage().getContent()));
String response = builder.toString();
return response;
}
private List<ChatCompletionRequest> getChatCompletionRequest(Prompt prompt) {
List<ChatMessage> chatMessages = convertToChatMessages(prompt.getMessages());
List<ChatCompletionRequest> chatCompletionRequests = new ArrayList<>();
ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
.model(this.model)
.temperature(this.temperature)
.messages(chatMessages)
.build();
chatCompletionRequests.add(chatCompletionRequest);
return chatCompletionRequests;
}
private List<ChatMessage> convertToChatMessages(List<Message> messages) {
List<ChatMessage> chatMessages = new ArrayList<>();
for (Message promptMessage : messages) {
MessageType promptMessageType = promptMessage.getMessageType();
switch (promptMessageType) {
case USER, ASSISTANT, SYSTEM -> chatMessages.add(newChatMessage(promptMessage));
case FUNCTION -> logger.error(
"Cannot send a Spring AI Function MessageType to the ChatGPT API; use 'system', 'user' or 'ai' message types.");
default ->
logger.error("Unknown Spring AI Chat MessageType; use 'system', 'human' or 'ai' message types.");
}
var chatCompletion = completionEntity.getBody();
if (chatCompletion == null) {
logger.warn("No chat completion returned for request: {}", chatCompletionMessages);
return new AiResponse(List.of());
}
return chatMessages;
RateLimit rateLimits = OpenAiResponseHeaderExtractor.extractAiResponseHeaders(completionEntity);
List<Generation> generations = chatCompletion.choices().stream().map(choice -> {
return new Generation(choice.message().content(), Map.of("role", choice.message().role().name()))
.withChoiceMetadata(ChoiceMetadata.from(choice.finishReason().name(), null));
}).toList();
return new AiResponse(generations,
OpenAiGenerationMetadata.from(completionEntity.getBody()).withRateLimit(rateLimits));
}
private ChatMessage newChatMessage(Message message) {
return new ChatMessage(message.getMessageType().getValue(), message.getContent());
@Override
public Flux<AiResponse> generateStream(Prompt prompt) {
List<Message> messages = prompt.getMessages();
List<ChatCompletionMessage> chatCompletionMessages = messages.stream()
.map(m -> new ChatCompletionMessage(m.getContent(),
ChatCompletionMessage.Role.valueOf(m.getMessageType().getValue())))
.toList();
Flux<OpenAiApi.ChatCompletionChunk> completionChunks = this.openAiApi
.chatCompletionStream(new OpenAiApi.ChatCompletionRequest(chatCompletionMessages, this.model,
this.temperature.floatValue(), true));
// For chunked responses, only the first chunk contains the choice role.
// The rest of the chunks with same ID share the same role.
ConcurrentHashMap<String, String> roleMap = new ConcurrentHashMap<>();
// An alternative implementation that returns Flux<Generation> instead of
// Flux<AiResponse>.
// Flux<Generation> generationFlux = completionChunks.map(chunk -> {
// String chunkId = chunk.id();
// return chunk.choices().stream()
// .map(choice -> {
// if (choice.delta().role() != null) {
// roleMap.putIfAbsent(chunkId, choice.delta().role().name());
// }
// return new Generation(choice.delta().content(),
// Map.of("role", roleMap.get(chunkId)));
// })
// .toList();
// }).flatMapIterable(generations -> generations);
// return generationFlux;
return completionChunks.map(chunk -> {
String chunkId = chunk.id();
List<Generation> generations = chunk.choices().stream().map(choice -> {
if (choice.delta().role() != null) {
roleMap.putIfAbsent(chunkId, choice.delta().role().name());
}
var generation = new Generation(choice.delta().content(), Map.of("role", roleMap.get(chunkId)));
if (choice.finishReason() != null) {
generation = generation.withChoiceMetadata(ChoiceMetadata.from(choice.finishReason().name(), null));
}
return generation;
}).toList();
return new AiResponse(generations);
});
}
}

34
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiSseResponse.java
View File

@@ -1,34 +0,0 @@
package org.springframework.ai.openai.client;
import com.fasterxml.jackson.annotation.JsonProperty;
import java.util.List;
public record OpenAiSseResponse(
@JsonProperty("created") Integer created,
@JsonProperty("model") String model,
@JsonProperty("id") String id,
@JsonProperty("system_fingerprint") String systemFingerprint,
@JsonProperty("choices") List<Choice> choices,
@JsonProperty("object") String object) {
public record Choice(
@JsonProperty("finish_reason") String finishReason,
@JsonProperty("delta") Delta delta,
@JsonProperty("index") Integer index) {
public record Delta(
@JsonProperty("role") String role,
@JsonProperty("content") String content) {
}
}
}

107
spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiStreamClient.java
View File

@@ -1,107 +0,0 @@
/*
* 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.openai.client;
import java.util.List;
import java.util.function.Predicate;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import org.springframework.ai.prompt.Prompt;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.reactive.function.client.WebClient;
public class OpenAiStreamClient implements AiStreamClient {
private Double temperature = 0.7;
private String model = "gpt-3.5-turbo";
private final Logger logger = LoggerFactory.getLogger(getClass());
private final WebClient webClient;
private final ObjectMapper objectMapper;
private final ParameterizedTypeReference<ServerSentEvent<String>> sseType;
public OpenAiStreamClient(String openAiApiToken) {
this("https://api.openai.com/", openAiApiToken);
}
public OpenAiStreamClient(String openAiEndpoint, String openAiApiToken) {
this.webClient = WebClient.builder()
.baseUrl(openAiEndpoint)
.defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
.defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + openAiApiToken)
.build();
this.objectMapper = new ObjectMapper().setSerializationInclusion(JsonInclude.Include.NON_NULL)
.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
.configure(DeserializationFeature.ACCEPT_EMPTY_STRING_AS_NULL_OBJECT, true);
this.sseType = new ParameterizedTypeReference<>() {
};
}
private WebClient.ResponseSpec createChatCompletion(ChatCompletionsRequest chatCompletionsRequest) {
return this.webClient.post()
.uri("/v1/chat/completions")
.bodyValue(objectMapper.convertValue(chatCompletionsRequest, JsonNode.class))
.retrieve();
}
@Override
public Flux<OpenAiSseResponse> generateStream(Prompt prompt) {
List<OpenAiChatMessage> openAiChatMessages = prompt.getMessages()
.stream()
.map(message -> new OpenAiChatMessage.Builder().role(message.getMessageTypeValue())
.content(message.getContent())
.build())
.toList();
ChatCompletionsRequest chatCompletionsRequest = new ChatCompletionsRequest.Builder().stream(true)
.model(this.model)
.temperature(this.temperature)
.messages(openAiChatMessages)
.build();
logger.trace("ChatMessages: {}", chatCompletionsRequest.getMessages());
return createChatCompletion(chatCompletionsRequest).bodyToFlux(sseType)
.map(ServerSentEvent::data)
.filter(Predicate.not("[DONE]"::equals))
.handle((data, sink) -> {
try {
sink.next(objectMapper.readValue(data, OpenAiSseResponse.class));
}
catch (JsonProcessingException e) {
sink.error(new RuntimeException(e));
}
});
}
}

141
spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java
View File

@@ -1,13 +1,25 @@
/*
* Copyright 2023-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.openai.embedding;
import java.util.ArrayList;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.Usage;
import com.theokanning.openai.embedding.EmbeddingRequest;
import com.theokanning.openai.service.OpenAiService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -16,91 +28,112 @@ import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.AbstractEmbeddingClient;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest;
import org.springframework.ai.openai.api.OpenAiApi.OpenAiApiException;
import org.springframework.ai.openai.api.OpenAiApi.Usage;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
/**
* Open AI Embedding Client implementation.
*
* @author Christian Tzolov
*/
public class OpenAiEmbeddingClient extends AbstractEmbeddingClient {
private static final Logger logger = LoggerFactory.getLogger(OpenAiEmbeddingClient.class);
private final OpenAiService openAiService;
public static final String DEFAULT_OPENAI_EMBEDDING_MODEL = "text-embedding-ada-002";
private final String model;
public final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(OpenAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.build();
private final OpenAiApi openAiApi;
private final String embeddingModelName;
private final MetadataMode metadataMode;
public OpenAiEmbeddingClient(OpenAiService openAiService) {
this(openAiService, "text-embedding-ada-002");
public OpenAiEmbeddingClient(OpenAiApi openAiApi) {
this(openAiApi, DEFAULT_OPENAI_EMBEDDING_MODEL);
}
public OpenAiEmbeddingClient(OpenAiService openAiService, String model) {
this(openAiService, model, MetadataMode.EMBED);
public OpenAiEmbeddingClient(OpenAiApi openAiApi, String embeddingModel) {
this(openAiApi, embeddingModel, MetadataMode.EMBED);
}
public OpenAiEmbeddingClient(OpenAiService openAiService, String model, MetadataMode metadataMode) {
Assert.notNull(openAiService, "OpenAiService must not be null");
public OpenAiEmbeddingClient(OpenAiApi openAiApi, String model, MetadataMode metadataMode) {
Assert.notNull(openAiApi, "OpenAiService must not be null");
Assert.notNull(model, "Model must not be null");
Assert.notNull(metadataMode, "metadataMode must not be null");
this.openAiService = openAiService;
this.model = model;
this.openAiApi = openAiApi;
this.embeddingModelName = model;
this.metadataMode = metadataMode;
}
@Override
public List<Double> embed(String text) {
EmbeddingRequest embeddingRequest = EmbeddingRequest.builder().input(List.of(text)).model(this.model).build();
com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult = this.openAiService
.createEmbeddings(embeddingRequest);
return generateEmbeddingResponse(nativeEmbeddingResult).getData().get(0).getEmbedding();
}
public List<Double> embed(Document document) {
EmbeddingRequest embeddingRequest = EmbeddingRequest.builder()
.input(List.of(document.getFormattedContent(this.metadataMode)))
.model(this.model)
.build();
com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult = this.openAiService
.createEmbeddings(embeddingRequest);
return generateEmbeddingResponse(nativeEmbeddingResult).getData().get(0).getEmbedding();
Assert.notNull(document, "Document must not be null");
return this.embed(document.getFormattedContent(this.metadataMode));
}
@Override
public List<Double> embed(String text) {
Assert.notNull(text, "Text must not be null");
return this.embed(List.of(text)).iterator().next();
}
@Override
public List<List<Double>> embed(List<String> texts) {
EmbeddingResponse embeddingResponse = embedForResponse(texts);
return embeddingResponse.getData().stream().map(emb -> emb.getEmbedding()).toList();
Assert.notNull(texts, "Texts must not be null");
EmbeddingRequest<List<String>> request = new EmbeddingRequest<>(texts, this.embeddingModelName);
return this.retryTemplate.execute(ctx -> {
EmbeddingList<OpenAiApi.Embedding> body = this.openAiApi.embeddings(request).getBody();
if (body == null) {
logger.warn("No embeddings returned for request: {}", request);
return List.of();
}
return body.data().stream().map(embedding -> embedding.embedding()).toList();
});
}
@Override
public EmbeddingResponse embedForResponse(List<String> texts) {
EmbeddingRequest embeddingRequest = EmbeddingRequest.builder().input(texts).model(this.model).build();
com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult = this.openAiService
.createEmbeddings(embeddingRequest);
return generateEmbeddingResponse(nativeEmbeddingResult);
}
private EmbeddingResponse generateEmbeddingResponse(
com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult) {
List<Embedding> data = generateEmbeddingList(nativeEmbeddingResult.getData());
Map<String, Object> metadata = generateMetadata(nativeEmbeddingResult.getModel(),
nativeEmbeddingResult.getUsage());
return new EmbeddingResponse(data, metadata);
}
Assert.notNull(texts, "Texts must not be null");
private List<Embedding> generateEmbeddingList(List<com.theokanning.openai.embedding.Embedding> nativeData) {
List<Embedding> data = new ArrayList<>();
for (com.theokanning.openai.embedding.Embedding nativeDatum : nativeData) {
List<Double> nativeDatumEmbedding = nativeDatum.getEmbedding();
int nativeIndex = nativeDatum.getIndex();
Embedding embedding = new Embedding(nativeDatumEmbedding, nativeIndex);
data.add(embedding);
}
return data;
return this.retryTemplate.execute(ctx -> {
EmbeddingRequest<List<String>> request = new EmbeddingRequest<>(texts, this.embeddingModelName);
EmbeddingList<OpenAiApi.Embedding> embeddingResponse = this.openAiApi.embeddings(request).getBody();
if (embeddingResponse == null) {
logger.warn("No embeddings returned for request: {}", request);
return new EmbeddingResponse(List.of(), Map.of());
}
Map<String, Object> metadata = generateMetadata(embeddingResponse.model(), embeddingResponse.usage());
List<Embedding> embeddings = embeddingResponse.data()
.stream()
.map(e -> new Embedding(e.embedding(), e.index()))
.toList();
return new EmbeddingResponse(embeddings, metadata);
});
}
private Map<String, Object> generateMetadata(String model, Usage usage) {
Map<String, Object> metadata = new HashMap<>();
metadata.put("model", model);
metadata.put("prompt-tokens", usage.getPromptTokens());
metadata.put("completion-tokens", usage.getCompletionTokens());
metadata.put("total-tokens", usage.getTotalTokens());
metadata.put("prompt-tokens", usage.promptTokens());
metadata.put("completion-tokens", usage.completionTokens());
metadata.put("total-tokens", usage.totalTokens());
return metadata;
}

9
spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiGenerationMetadata.java
View File

@@ -16,11 +16,10 @@
package org.springframework.ai.openai.metadata;
import com.theokanning.openai.completion.chat.ChatCompletionResult;
import org.springframework.ai.metadata.GenerationMetadata;
import org.springframework.ai.metadata.RateLimit;
import org.springframework.ai.metadata.Usage;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.metadata.support.OpenAiHttpResponseHeadersInterceptor;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -38,10 +37,10 @@ public class OpenAiGenerationMetadata implements GenerationMetadata {
protected static final String AI_METADATA_STRING = "{ @type: %1$s, id: %2$s, usage: %3$s, rateLimit: %4$s }";
public static OpenAiGenerationMetadata from(ChatCompletionResult result) {
public static OpenAiGenerationMetadata from(OpenAiApi.ChatCompletion result) {
Assert.notNull(result, "OpenAI ChatCompletionResult must not be null");
OpenAiUsage usage = OpenAiUsage.from(result.getUsage());
OpenAiGenerationMetadata generationMetadata = new OpenAiGenerationMetadata(result.getId(), usage);
OpenAiUsage usage = OpenAiUsage.from(result.usage());
OpenAiGenerationMetadata generationMetadata = new OpenAiGenerationMetadata(result.id(), usage);
OpenAiHttpResponseHeadersInterceptor.applyTo(generationMetadata);
return generationMetadata;
}

15
spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java
View File

@@ -17,6 +17,7 @@
package org.springframework.ai.openai.metadata;
import org.springframework.ai.metadata.Usage;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.util.Assert;
/**
@@ -30,34 +31,34 @@ import org.springframework.util.Assert;
*/
public class OpenAiUsage implements Usage {
public static OpenAiUsage from(com.theokanning.openai.Usage usage) {
public static OpenAiUsage from(OpenAiApi.Usage usage) {
return new OpenAiUsage(usage);
}
private final com.theokanning.openai.Usage usage;
private final OpenAiApi.Usage usage;
protected OpenAiUsage(com.theokanning.openai.Usage usage) {
protected OpenAiUsage(OpenAiApi.Usage usage) {
Assert.notNull(usage, "OpenAI Usage must not be null");
this.usage = usage;
}
protected com.theokanning.openai.Usage getUsage() {
protected OpenAiApi.Usage getUsage() {
return this.usage;
}
@Override
public Long getPromptTokens() {
return getUsage().getPromptTokens();
return getUsage().promptTokens().longValue();
}
@Override
public Long getGenerationTokens() {
return getUsage().getCompletionTokens();
return getUsage().completionTokens().longValue();
}
@Override
public Long getTotalTokens() {
return getUsage().getTotalTokens();
return getUsage().totalTokens().longValue();
}
@Override

198
spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/support/OpenAiResponseHeaderExtractor.java Normal file
View File

@@ -0,0 +1,198 @@
/*
* 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.openai.metadata.support;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.Arrays;
import java.util.function.Predicate;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.metadata.RateLimit;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion;
import org.springframework.ai.openai.metadata.OpenAiRateLimit;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.REQUESTS_LIMIT_HEADER;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.REQUESTS_REMAINING_HEADER;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.REQUESTS_RESET_HEADER;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.TOKENS_LIMIT_HEADER;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.TOKENS_REMAINING_HEADER;
import static org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders.TOKENS_RESET_HEADER;
/**
* Utility used to extract known HTTP response headers for the {@literal OpenAI} API.
*
* @author John Blum
* @author Christian Tzolov
* @since 0.7.0
*/
public class OpenAiResponseHeaderExtractor {
private static final Logger logger = LoggerFactory.getLogger(OpenAiResponseHeaderExtractor.class);
public static RateLimit extractAiResponseHeaders(ResponseEntity<ChatCompletion> response) {
Long requestsLimit = getHeaderAsLong(response, REQUESTS_LIMIT_HEADER.getName());
Long requestsRemaining = getHeaderAsLong(response, REQUESTS_REMAINING_HEADER.getName());
Long tokensLimit = getHeaderAsLong(response, TOKENS_LIMIT_HEADER.getName());
Long tokensRemaining = getHeaderAsLong(response, TOKENS_REMAINING_HEADER.getName());
Duration requestsReset = getHeaderAsDuration(response, REQUESTS_RESET_HEADER.getName());
Duration tokensReset = getHeaderAsDuration(response, TOKENS_RESET_HEADER.getName());
return new OpenAiRateLimit(requestsLimit, requestsRemaining, requestsReset, tokensLimit, tokensRemaining,
tokensReset);
}
private static Duration getHeaderAsDuration(ResponseEntity<ChatCompletion> response, String headerName) {
var headers = response.getHeaders();
if (headers.containsKey(headerName)) {
var values = headers.get(headerName);
if (!CollectionUtils.isEmpty(values)) {
return DurationFormatter.TIME_UNIT.parse(values.get(0));
}
}
return null;
}
private static Long getHeaderAsLong(ResponseEntity<ChatCompletion> response, String headerName) {
var headers = response.getHeaders();
if (headers.containsKey(headerName)) {
var values = headers.get(headerName);
if (!CollectionUtils.isEmpty(values)) {
return parseLong(headerName, values.get(0));
}
}
return null;
}
private static Long parseLong(String headerName, String headerValue) {
if (StringUtils.hasText(headerValue)) {
try {
return Long.parseLong(headerValue.trim());
}
catch (NumberFormatException e) {
logger.warn("Value [{}] for HTTP header [{}] is not valid: {}", headerName, headerValue,
e.getMessage());
}
}
return null;
}
enum DurationFormatter {
TIME_UNIT("\\d+[a-zA-Z]{1,2}");
private final Pattern pattern;
DurationFormatter(String durationPattern) {
this.pattern = Pattern.compile(durationPattern);
}
public Duration parse(String text) {
Assert.hasText(text, "Text [%s] to parse as a Duration must not be null or empty".formatted(text));
Matcher matcher = this.pattern.matcher(text);
Duration total = Duration.ZERO;
while (matcher.find()) {
String value = matcher.group();
total = total.plus(Unit.parseUnit(value).toDuration(value));
}
return total;
}
enum Unit {
NANOSECONDS("ns", "nanoseconds", ChronoUnit.NANOS), MICROSECONDS("us", "microseconds", ChronoUnit.MICROS),
MILLISECONDS("ms", "milliseconds", ChronoUnit.MILLIS), SECONDS("s", "seconds", ChronoUnit.SECONDS),
MINUTES("m", "minutes", ChronoUnit.MINUTES), HOURS("h", "hours", ChronoUnit.HOURS),
DAYS("d", "days", ChronoUnit.DAYS);
private final String name;
private final String symbol;
private final ChronoUnit unit;
Unit(String symbol, String name, ChronoUnit unit) {
this.symbol = symbol;
this.name = name;
this.unit = unit;
}
static Unit parseUnit(String value) {
String symbol = parseSymbol(value);
return Arrays.stream(values())
.filter(unit -> unit.getSymbol().equalsIgnoreCase(symbol))
.findFirst()
.orElseThrow(() -> new IllegalStateException(
"Value [%s] does not contain a valid time unit".formatted(value)));
}
private static String parse(String value, Predicate<Character> predicate) {
Assert.hasText(value, "Value [%s] must not be null or empty".formatted(value));
StringBuilder builder = new StringBuilder();
for (char character : value.toCharArray()) {
if (predicate.test(character)) {
builder.append(character);
}
}
return builder.toString();
}
private static String parseSymbol(String value) {
return parse(value, Character::isLetter);
}
private static Long parseTime(String value) {
return Long.parseLong(parse(value, Character::isDigit));
}
public String getName() {
return this.name;
}
public String getSymbol() {
return this.symbol;
}
public ChronoUnit getUnit() {
return this.unit;
}
public Duration toDuration(String value) {
return Duration.of(parseTime(value), getUnit());
}
}
}
}

96
spring-ai-openai/src/test/java/org/springframework/ai/openai/MockOpenAiTestConfiguration.java
View File

@@ -1,96 +0,0 @@
/*
* 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.openai;
import static org.springframework.ai.test.config.MockAiTestConfiguration.SPRING_AI_API_PATH;
import java.time.Duration;
import java.util.UUID;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.ai.openai.metadata.support.OpenAiHttpResponseHeadersInterceptor;
import org.springframework.ai.test.config.MockAiTestConfiguration;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Import;
import org.springframework.context.annotation.Profile;
import org.springframework.test.web.servlet.MockMvc;
import okhttp3.HttpUrl;
import okhttp3.OkHttpClient;
import okhttp3.mockwebserver.Dispatcher;
import okhttp3.mockwebserver.MockWebServer;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
/**
* {@link SpringBootConfiguration} for testing {@literal OpenAI's} API using mock objects.
* <p>
* This test configuration allows Spring AI framework developers to mock OpenAI's API with
* Spring {@link MockMvc} and a test provided Spring Web MVC
* {@link org.springframework.web.bind.annotation.RestController}.
* <p>
* This test configuration makes use of the OkHttp3 {@link MockWebServer} and
* {@link Dispatcher} to integrate with Spring {@link MockMvc}.
*
* @author John Blum
* @see org.springframework.boot.SpringBootConfiguration
* @see org.springframework.ai.test.config.MockAiTestConfiguration
* @since 0.7.0
*/
@SpringBootConfiguration
@Profile("spring-ai-openai-mocks")
@Import(MockAiTestConfiguration.class)
@SuppressWarnings("unused")
public class MockOpenAiTestConfiguration {
@Bean
OpenAiService theoOpenAiService(MockWebServer webServer) {
String apiKey = UUID.randomUUID().toString();
Duration timeout = Duration.ofSeconds(60);
ObjectMapper objectMapper = OpenAiService.defaultObjectMapper();
OkHttpClient httpClient = new OkHttpClient.Builder(OpenAiService.defaultClient(apiKey, timeout))
.addInterceptor(new OpenAiHttpResponseHeadersInterceptor())
.build();
HttpUrl baseUrl = webServer.url(SPRING_AI_API_PATH.concat("/"));
Retrofit retrofit = new Retrofit.Builder().baseUrl(baseUrl)
.addConverterFactory(JacksonConverterFactory.create(objectMapper))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.client(httpClient)
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiService(api);
}
@Bean
OpenAiClient apiClient(OpenAiService openAiService) {
return new OpenAiClient(openAiService);
}
}

23
spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java
View File

@@ -1,12 +1,8 @@
package org.springframework.ai.openai;
import java.time.Duration;
import com.theokanning.openai.service.OpenAiService;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.ai.openai.client.OpenAiStreamClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.context.annotation.Bean;
@@ -16,9 +12,9 @@ import org.springframework.util.StringUtils;
public class OpenAiTestConfiguration {
@Bean
public OpenAiService theoOpenAiService() {
public OpenAiApi openAiApi() {
String apiKey = getApiKey();
OpenAiService openAiService = new OpenAiService(apiKey, Duration.ofSeconds(60));
OpenAiApi openAiService = new OpenAiApi(apiKey);
return openAiService;
}
@@ -32,20 +28,15 @@ public class OpenAiTestConfiguration {
}
@Bean
public OpenAiClient openAiClient(OpenAiService theoOpenAiService) {
OpenAiClient openAiClient = new OpenAiClient(theoOpenAiService);
public OpenAiClient openAiClient(OpenAiApi api) {
OpenAiClient openAiClient = new OpenAiClient(api);
openAiClient.setTemperature(0.3);
return openAiClient;
}
@Bean
public EmbeddingClient openAiEmbeddingClient(OpenAiService theoOpenAiService) {
return new OpenAiEmbeddingClient(theoOpenAiService);
}
@Bean
public OpenAiStreamClient openAiStreamClient() {
return new OpenAiStreamClient(getApiKey());
public EmbeddingClient openAiEmbeddingClient(OpenAiApi api) {
return new OpenAiEmbeddingClient(api);
}
}

9
spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java
View File

@@ -140,11 +140,12 @@ class OpenAiClientIT extends AbstractIT {
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = openAiStreamClient.generateStream(prompt)
.map(OpenAiSseResponse::choices)
.toStream()
.collectList()
.block()
.stream()
.map(AiResponse::getGenerations)
.flatMap(List::stream)
.map(OpenAiSseResponse.Choice::delta)
.map(OpenAiSseResponse.Choice.Delta::content)
.map(Generation::getContent)
.collect(Collectors.joining());
ActorsFilmsRecord actorsFilms = outputParser.parse(generationTextFromStream);

155
spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadataTests.java → spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadata2Tests.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023 the original author or authors.
* Copyright 2023-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.
@@ -16,12 +16,9 @@
package org.springframework.ai.openai.client;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.test.config.MockAiTestConfiguration.SPRING_AI_API_PATH;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.springframework.ai.client.AiResponse;
@@ -30,49 +27,54 @@ import org.springframework.ai.metadata.GenerationMetadata;
import org.springframework.ai.metadata.PromptMetadata;
import org.springframework.ai.metadata.RateLimit;
import org.springframework.ai.metadata.Usage;
import org.springframework.ai.openai.MockOpenAiTestConfiguration;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders;
import org.springframework.ai.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.autoconfigure.web.client.RestClientTest;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Import;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.test.context.ActiveProfiles;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.web.servlet.MockMvc;
import org.springframework.test.web.servlet.setup.MockMvcBuilders;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.context.request.WebRequest;
import org.springframework.test.web.client.MockRestServiceServer;
import org.springframework.web.client.RestClient;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.test.web.client.match.MockRestRequestMatchers.header;
import static org.springframework.test.web.client.match.MockRestRequestMatchers.method;
import static org.springframework.test.web.client.match.MockRestRequestMatchers.requestTo;
import static org.springframework.test.web.client.response.MockRestResponseCreators.withSuccess;
/**
* Tests using the {@link OpenAiClient} to send an {@literal OpenAI} API request (chat
* completion) to test the presence of {@link GenerationMetadata} in the
* {@link AiResponse}.
*
* @author John Blum
* @author Christian Tzolov
* @since 0.7.0
*/
@SpringBootTest
@ActiveProfiles("spring-ai-openai-mocks")
@ContextConfiguration(classes = OpenAiClientWithGenerationMetadataTests.TestConfiguration.class)
@SuppressWarnings("unused")
class OpenAiClientWithGenerationMetadataTests {
@RestClientTest(OpenAiClientWithGenerationMetadata2Tests.Config.class)
public class OpenAiClientWithGenerationMetadata2Tests {
private static String TEST_API_KEY = "sk-1234567890";
@Autowired
private OpenAiClient aiClient;
private OpenAiClient openAiClient;
@Autowired
private MockRestServiceServer server;
@AfterEach
void resetMockServer() {
server.reset();
}
@Test
void aiResponseContainsAiMetadata() {
prepareMock();
Prompt prompt = new Prompt("Reach for the sky.");
AiResponse response = this.aiClient.generate(prompt);
AiResponse response = this.openAiClient.generate(prompt);
assertThat(response).isNotNull();
@@ -119,65 +121,58 @@ class OpenAiClientWithGenerationMetadataTests {
});
}
@SpringBootConfiguration
@Import(MockOpenAiTestConfiguration.class)
static class TestConfiguration {
private void prepareMock() {
@Bean
MockMvc mockMvc() {
return MockMvcBuilders.standaloneSetup(new SpringOpenAiChatCompletionsController()).build();
}
HttpHeaders httpHeaders = new HttpHeaders();
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_LIMIT_HEADER.getName(), "4000");
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_REMAINING_HEADER.getName(), "999");
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_RESET_HEADER.getName(), "2d16h15m29s");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_LIMIT_HEADER.getName(), "725000");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_REMAINING_HEADER.getName(), "112358");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_RESET_HEADER.getName(), "27h55s451ms");
server.expect(requestTo("/v1/chat/completions"))
.andExpect(method(HttpMethod.POST))
.andExpect(header(HttpHeaders.AUTHORIZATION, "Bearer " + TEST_API_KEY))
.andRespond(withSuccess(getJson(), MediaType.APPLICATION_JSON).headers(httpHeaders));
}
@RestController
@RequestMapping(SPRING_AI_API_PATH)
@SuppressWarnings("all")
static class SpringOpenAiChatCompletionsController {
private String getJson() {
return """
{
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-3.5-turbo-0613",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "I surrender!"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 9,
"completion_tokens": 12,
"total_tokens": 21
}
}
""";
}
@PostMapping("/v1/chat/completions")
ResponseEntity<?> chatCompletions(WebRequest request) {
@SpringBootConfiguration
static class Config {
String json = getJson();
ResponseEntity<?> response = ResponseEntity.status(HttpStatusCode.valueOf(200))
.contentType(MediaType.APPLICATION_JSON)
.contentLength(json.getBytes(StandardCharsets.UTF_8).length)
.headers(httpHeaders -> {
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_LIMIT_HEADER.getName(), "4000");
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_REMAINING_HEADER.getName(), "999");
httpHeaders.set(OpenAiApiResponseHeaders.REQUESTS_RESET_HEADER.getName(), "2d16h15m29s");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_LIMIT_HEADER.getName(), "725000");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_REMAINING_HEADER.getName(), "112358");
httpHeaders.set(OpenAiApiResponseHeaders.TOKENS_RESET_HEADER.getName(), "27h55s451ms");
})
.body(getJson());
return response;
@Bean
public OpenAiApi chatCompletionApi(RestClient.Builder builder) {
return new OpenAiApi("", TEST_API_KEY, builder);
}
private String getJson() {
return """
{
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-3.5-turbo-0613",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "I surrender!"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 9,
"completion_tokens": 12,
"total_tokens": 21
}
}
""";
@Bean
public OpenAiClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiClient(openAiApi);
}
}

21
spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java
View File

@@ -1,3 +1,18 @@
/*
* Copyright 2023-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.openai.embedding;
import org.junit.jupiter.api.Test;
@@ -20,13 +35,11 @@ class EmbeddingIT {
assertThat(embeddingClient).isNotNull();
EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World"));
System.out.println(embeddingResponse);
assertThat(embeddingResponse.getData()).hasSize(1);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getMetadata()).containsEntry("model", "text-embedding-ada-002-v2");
assertThat(embeddingResponse.getMetadata()).containsEntry("completion-tokens", 0L);
assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2L);
assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2L);
assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2);
assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2);
assertThat(embeddingClient.dimensions()).isEqualTo(1536);
}

11
spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java
View File

@@ -1,10 +1,14 @@
package org.springframework.ai.openai.testutils;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.openai.client.OpenAiStreamClient;
import org.springframework.ai.client.AiStreamClient;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.PromptTemplate;
import org.springframework.ai.prompt.messages.Message;
@@ -13,9 +17,6 @@ import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
@@ -27,7 +28,7 @@ public abstract class AbstractIT {
protected AiClient openAiClient;
@Autowired
protected OpenAiStreamClient openAiStreamClient;
protected AiStreamClient openAiStreamClient;
@Value("classpath:/prompts/eval/qa-evaluator-accurate-answer.st")
protected Resource qaEvaluatorAccurateAnswerResource;

12
spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java
View File

@@ -17,16 +17,15 @@
package org.springframework.ai.openai.transformer;
import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.ai.transformer.ContentFormatTransformer;
import org.springframework.ai.transformer.KeywordMetadataEnricher;
@@ -155,19 +154,18 @@ public class MetadataTransformerIT {
public static class OpenAiTestConfiguration {
@Bean
public OpenAiService theoOpenAiService() throws IOException {
public OpenAiApi openAiApi() throws IOException {
String apiKey = System.getenv("OPENAI_API_KEY");
if (!StringUtils.hasText(apiKey)) {
throw new IllegalArgumentException(
"You must provide an API key. Put it in an environment variable under the name OPENAI_API_KEY");
}
OpenAiService openAiService = new OpenAiService(apiKey, Duration.ofSeconds(60));
return openAiService;
return new OpenAiApi(apiKey);
}
@Bean
public OpenAiClient openAiClient(OpenAiService theoOpenAiService) {
OpenAiClient openAiClient = new OpenAiClient(theoOpenAiService);
public OpenAiClient openAiClient(OpenAiApi openAiApi) {
OpenAiClient openAiClient = new OpenAiClient(openAiApi);
openAiClient.setTemperature(0.3);
return openAiClient;
}

14
spring-ai-spring-boot-autoconfigure/pom.xml
View File

@@ -72,6 +72,14 @@
<optional>true</optional>
</dependency>
<!-- Because of Pinecone compatability issues downgrade
netty-codec-http2 from 4.1.101.Final to 4.1.100.Final -->
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-codec-http2</artifactId>
<version>4.1.100.Final</version>
</dependency>
<!-- Milvus Vector Store -->
<dependency>
<groupId>org.springframework.ai</groupId>
@@ -140,12 +148,6 @@
<optional>true</optional>
</dependency>
<dependency>
<groupId>com.squareup.retrofit2</groupId>
<artifactId>converter-jackson</artifactId>
<version>2.9.0</version>
</dependency>
<!-- test dependencies -->
<dependency>

70
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java
View File

@@ -16,18 +16,10 @@
package org.springframework.ai.autoconfigure.openai;
import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFIG_PREFIX;
import java.time.Duration;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.springframework.ai.autoconfigure.NativeHints;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.ai.openai.metadata.support.OpenAiHttpResponseHeadersInterceptor;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -35,27 +27,24 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ImportRuntimeHints;
import org.springframework.util.StringUtils;
import okhttp3.OkHttpClient;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.web.client.RestClient;
@AutoConfiguration
@ConditionalOnClass(OpenAiService.class)
@ConditionalOnClass(OpenAiApi.class)
@EnableConfigurationProperties(OpenAiProperties.class)
@ImportRuntimeHints(NativeHints.class)
public class OpenAiAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public OpenAiClient openAiClient(OpenAiProperties openAiProperties) {
public OpenAiApi openAiApi(OpenAiProperties openAiProperties) {
return new OpenAiApi(openAiProperties.getBaseUrl(), openAiProperties.getApiKey(), RestClient.builder());
}
OpenAiService openAiService = theoOpenAiService(openAiProperties, openAiProperties.getBaseUrl(),
openAiProperties.getApiKey(), openAiProperties.getDuration());
OpenAiClient openAiClient = new OpenAiClient(openAiService);
@Bean
@ConditionalOnMissingBean
public OpenAiClient openAiClient(OpenAiApi openAiApi, OpenAiProperties openAiProperties) {
OpenAiClient openAiClient = new OpenAiClient(openAiApi);
openAiClient.setTemperature(openAiProperties.getTemperature());
openAiClient.setModel(openAiProperties.getModel());
@@ -64,43 +53,8 @@ public class OpenAiAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public EmbeddingClient openAiEmbeddingClient(OpenAiProperties openAiProperties) {
OpenAiService openAiService = theoOpenAiService(openAiProperties, openAiProperties.getEmbedding().getBaseUrl(),
openAiProperties.getEmbedding().getApiKey(), openAiProperties.getDuration());
return new OpenAiEmbeddingClient(openAiService, openAiProperties.getEmbedding().getModel());
}
private OpenAiService theoOpenAiService(OpenAiProperties properties, String baseUrl, String apiKey,
Duration duration) {
if ("https://api.openai.com".equals(baseUrl) && !StringUtils.hasText(apiKey)) {
throw new IllegalArgumentException(
"You must provide an API key with the property name " + CONFIG_PREFIX + ".api-key");
}
ObjectMapper mapper = OpenAiService.defaultObjectMapper();
OkHttpClient.Builder clientBuilder = new OkHttpClient.Builder(OpenAiService.defaultClient(apiKey, duration));
if (properties.getMetadata().isRateLimitMetricsEnabled()) {
clientBuilder.addInterceptor(new OpenAiHttpResponseHeadersInterceptor());
}
OkHttpClient client = clientBuilder.build();
// Waiting for https://github.com/TheoKanning/openai-java/issues/249 to be
// resolved.
Retrofit retrofit = new Retrofit.Builder().baseUrl(baseUrl)
.client(client)
.addConverterFactory(JacksonConverterFactory.create(mapper))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiService(api);
public EmbeddingClient openAiEmbeddingClient(OpenAiApi openAiApi, OpenAiProperties openAiProperties) {
return new OpenAiEmbeddingClient(openAiApi, openAiProperties.getEmbedding().getModel());
}
}

16
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiProperties.java
View File

@@ -16,23 +16,17 @@
package org.springframework.ai.autoconfigure.openai;
import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFIG_PREFIX;
import java.time.Duration;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
@ConfigurationProperties(CONFIG_PREFIX)
@ConfigurationProperties(OpenAiProperties.CONFIG_PREFIX)
public class OpenAiProperties {
public static final String CONFIG_PREFIX = "spring.ai.openai";
private Double temperature = 0.7;
private Duration duration = Duration.ofSeconds(60);
private final Embedding embedding = new Embedding(this);
private final Metadata metadata = new Metadata();
@@ -75,14 +69,6 @@ public class OpenAiProperties {
this.temperature = temperature;
}
public Duration getDuration() {
return this.duration;
}
public void setDuration(Duration duration) {
this.duration = duration;
}
public Embedding getEmbedding() {
return this.embedding;
}

1
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiPropertiesTests.java
View File

@@ -50,7 +50,6 @@ class OpenAiPropertiesTests {
assertThat(this.openAiProperties.getModel()).isEqualTo("claudia-shiffer-5");
assertThat(this.openAiProperties.getBaseUrl()).isEqualTo("https://api.openai.spring.io/eieioh");
assertThat(this.openAiProperties.getTemperature()).isEqualTo(0.5d);
assertThat(this.openAiProperties.getDuration()).isEqualTo(Duration.ofSeconds(30L));
OpenAiProperties.Embedding embedding = this.openAiProperties.getEmbedding();

14
spring-ai-vertex-ai/pom.xml
View File

@@ -28,22 +28,10 @@
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>google-cloud-aiplatform</artifactId>
<version>3.32.0</version>
<exclusions>
<exclusion>
<groupId>commons-logging</groupId>
<artifactId>commons-logging</artifactId>
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
<version>6.1.1</version>
<version>${spring-framework.version}</version>
</dependency>

22
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java
View File

@@ -16,29 +16,24 @@
package org.springframework.ai.vectorstore;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.images.builder.Transferable;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
@@ -121,16 +116,7 @@ public class BasicAuthChromaWhereIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

22
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java
View File

@@ -16,30 +16,25 @@
package org.springframework.ai.vectorstore;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
@@ -228,16 +223,7 @@ public class ChromaVectorStoreIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

22
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java
View File

@@ -16,28 +16,23 @@
package org.springframework.ai.vectorstore;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.vectorsore.ChromaVectorStore;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
@@ -155,16 +150,7 @@ public class TokenSecuredChromaWhereIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

3
vector-stores/spring-ai-milvus-store/pom.xml
View File

@@ -57,8 +57,9 @@
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai</artifactId>
<!-- <artifactId>spring-ai-openai</artifactId> -->
<!-- <artifactId>transformers-embedding</artifactId> -->
<artifactId>spring-ai-vertex-ai</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>

28
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
View File

@@ -25,8 +25,6 @@ import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
@@ -39,14 +37,12 @@ import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
import org.springframework.ai.vertex.api.VertexAiApi;
import org.springframework.ai.vertex.embedding.VertexAiEmbeddingClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -62,7 +58,8 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
// @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "PALM_API_KEY", matches = ".+")
public class MilvusVectorStoreIT {
private static DockerComposeContainer milvusContainer;
@@ -153,6 +150,8 @@ public class MilvusVectorStoreIT {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
@@ -246,7 +245,7 @@ public class MilvusVectorStoreIT {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "COSINE", "L2", "IP" })
@ValueSource(strings = { "COSINE", "IP" })
public void searchWithThreshold(String metricType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
@@ -308,16 +307,9 @@ public class MilvusVectorStoreIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new VertexAiEmbeddingClient(new VertexAiApi(System.getenv("PALM_API_KEY")));
// return new OpenAiEmbeddingClient(new
// OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

18
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java
View File

@@ -1,12 +1,9 @@
package org.springframework.ai.vectorstore;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -17,12 +14,10 @@ import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -210,16 +205,7 @@ class Neo4jVectorStoreIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

18
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -18,7 +18,6 @@ package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
@@ -27,8 +26,6 @@ import java.util.UUID;
import javax.sql.DataSource;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.Assert;
import org.junit.jupiter.api.Assertions;
@@ -38,12 +35,10 @@ import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
@@ -337,16 +332,7 @@ public class PgVectorStoreIT {
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
}

125
vector-stores/spring-ai-pinecone/README.md Normal file
View File

@@ -0,0 +1,125 @@
# Pinecone Vector Store
This readme walks you through setting up the Pinecone `VectorStore` to store document embeddings and perform similarity searches.
## What is Pinecone?
[Pinecone](https://www.pinecone.io/) is a popular cloud-based vector database, which allows you to store and search vectors efficiently.
## Prerequisites
1. Pinecone Account: Before you start, sign up for a [Pinecone account](https://app.pinecone.io/).
2. Pinecone Project: Once registered, create a new project, an index, and generate an API key. You'll need these details for configuration.
3. OpenAI Account: Create an account at [OpenAI Signup](https://platform.openai.com/signup) and generate the token at [API Keys](https://platform.openai.com/account/api-keys)
## Configuration
To set up `PineconeVectorStore`, gather the following details from your Pinecone account:
* Pinecone API Key
* Pinecone Environment
* Pinecone Project ID
* Pinecone Index Name
* Pinecone Namespace
> **Note**
> This information is available to you in the Pinecone UI portal.
When setting up embeddings, select a vector dimension of `1536`. This matches the dimensionality of OpenAI's model `text-embedding-ada-002`, which we'll be using for this guide.
Additionally, you'll need to provide your OpenAI API Key. Set it as an environment variable like so:
```bash
export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key'
```
## Repository
To acquire Spring AI artifacts, declare the Spring Snapshot repository:
```xml
<repository>
<id>spring-snapshots</id>
<name>Spring Snapshots</name>
<url>https://repo.spring.io/snapshot</url>
<releases>
<enabled>false</enabled>
</releases>
</repository>
```
## Dependencies
Add these dependencies to your project:
1. OpenAI: Required for calculating embeddings.
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
2. Pinecone
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-pinecone</artifactId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
## Sample Code
To configure Pinecone in your application, you can use the following setup:
```java
@Bean
public PineconeVectorStoreConfig pineconeVectorStoreConfig() {
return PineconeVectorStoreConfig.builder()
.withApiKey(<PINECONE_API_KEY>)
.withEnvironment("gcp-starter")
.withProjectId("89309e6")
.withIndexName("spring-ai-test-index")
.withNamespace("") // the free tier doesn't support namespaces.
.build();
}
```
Integrate with OpenAI's embeddings by adding the Spring Boot OpenAI starter to your project.
This provides you with an implementation of the Embeddings client:
```java
@Bean
public VectorStore vectorStore(PineconeVectorStoreConfig config, EmbeddingClient embeddingClient) {
return new PineconeVectorStore(config, embeddingClient);
}
```
In your main code, create some documents:
```java
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
new Document("The World is Big and Salvation Lurks Around the Corner"),
new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
```
Add the documents to Pinecone:
```java
vectorStore.add(List.of(document));
```
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
```
If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".

47
vector-stores/spring-ai-pinecone/pom.xml
View File

@@ -36,8 +36,31 @@
<groupId>io.pinecone</groupId>
<artifactId>pinecone-client</artifactId>
<version>${pinecone.version}</version>
<exclusions>
<exclusion>
<groupId>io.netty</groupId>
<artifactId>netty-codec</artifactId>
</exclusion>
<exclusion>
<groupId>io.netty</groupId>
<artifactId>netty-codec-http2</artifactId>
</exclusion>
</exclusions>
</dependency>
<!-- Overrides the pinecone's dependencies. TODO review after pincone verions is updated to > 0.6.0 -->
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-codec</artifactId>
<version>4.1.101.Final</version>
</dependency>
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-codec-http2</artifactId>
<version>4.1.100.Final</version>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java-util</artifactId>
@@ -45,19 +68,19 @@
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>