From e30be94a564fe751bf2a2e316925d8cb12ba2224 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Wed, 13 Dec 2023 15:48:28 +0100 Subject: [PATCH] Replace the theokanning with custom OpenAi API client - 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 - 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 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 --- pom.xml | 2 +- .../AzureOpenAiClientMetadataTests.java | 3 +- spring-ai-core/pom.xml | 2 +- .../ai}/client/AiStreamClient.java | 11 +- spring-ai-openai/pom.xml | 12 +- .../ai/openai/api/OpenAiApi.java | 718 ++++++++++++++++++ .../openai/client/ChatCompletionResponse.java | 31 - .../openai/client/ChatCompletionsRequest.java | 286 ------- .../ai/openai/client/OpenAiChatMessage.java | 91 --- .../ai/openai/client/OpenAiClient.java | 232 +++--- .../ai/openai/client/OpenAiSseResponse.java | 34 - .../ai/openai/client/OpenAiStreamClient.java | 107 --- .../embedding/OpenAiEmbeddingClient.java | 141 ++-- .../metadata/OpenAiGenerationMetadata.java | 9 +- .../ai/openai/metadata/OpenAiUsage.java | 15 +- .../OpenAiResponseHeaderExtractor.java | 198 +++++ .../openai/MockOpenAiTestConfiguration.java | 96 --- .../ai/openai/OpenAiTestConfiguration.java | 23 +- .../ai/openai/client/OpenAiClientIT.java | 9 +- ...AiClientWithGenerationMetadata2Tests.java} | 155 ++-- .../ai/openai/embedding/EmbeddingIT.java | 21 +- .../ai/openai/testutils/AbstractIT.java | 11 +- .../transformer/MetadataTransformerIT.java | 12 +- spring-ai-spring-boot-autoconfigure/pom.xml | 14 +- .../openai/OpenAiAutoConfiguration.java | 70 +- .../openai/OpenAiProperties.java | 16 +- .../openai/OpenAiPropertiesTests.java | 1 - spring-ai-vertex-ai/pom.xml | 14 +- .../vectorstore/BasicAuthChromaWhereIT.java | 22 +- .../ai/vectorstore/ChromaVectorStoreIT.java | 22 +- .../TokenSecuredChromaWhereIT.java | 22 +- vector-stores/spring-ai-milvus-store/pom.xml | 3 +- .../ai/vectorstore/MilvusVectorStoreIT.java | 28 +- .../ai/vectorstore/Neo4jVectorStoreIT.java | 18 +- .../ai/vectorstore/PgVectorStoreIT.java | 18 +- vector-stores/spring-ai-pinecone/README.md | 125 +++ vector-stores/spring-ai-pinecone/pom.xml | 47 +- 37 files changed, 1449 insertions(+), 1190 deletions(-) rename {spring-ai-openai/src/main/java/org/springframework/ai/openai => spring-ai-core/src/main/java/org/springframework/ai}/client/AiStreamClient.java (69%) create mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java delete mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionResponse.java delete mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionsRequest.java delete mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiChatMessage.java delete mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiSseResponse.java delete mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiStreamClient.java create mode 100644 spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/support/OpenAiResponseHeaderExtractor.java delete mode 100644 spring-ai-openai/src/test/java/org/springframework/ai/openai/MockOpenAiTestConfiguration.java rename spring-ai-openai/src/test/java/org/springframework/ai/openai/client/{OpenAiClientWithGenerationMetadataTests.java => OpenAiClientWithGenerationMetadata2Tests.java} (50%) create mode 100644 vector-stores/spring-ai-pinecone/README.md diff --git a/pom.xml b/pom.xml index 6666cdfe1..54939de61 100644 --- a/pom.xml +++ b/pom.xml @@ -84,8 +84,8 @@ 3.2.0 + 6.1.1 4.0.2 - 0.16.0 1.0.0-beta.3 0.6.1 4.31.1 diff --git a/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/client/AzureOpenAiClientMetadataTests.java b/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/client/AzureOpenAiClientMetadataTests.java index 8163461c4..2a330859e 100644 --- a/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/client/AzureOpenAiClientMetadataTests.java +++ b/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/client/AzureOpenAiClientMetadataTests.java @@ -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 { diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml index 937836404..6cd1c4617 100644 --- a/spring-ai-core/pom.xml +++ b/spring-ai-core/pom.xml @@ -42,7 +42,7 @@ org.springframework spring-webflux - 6.1.1 + ${spring-framework.version} diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/AiStreamClient.java b/spring-ai-core/src/main/java/org/springframework/ai/client/AiStreamClient.java similarity index 69% rename from spring-ai-openai/src/main/java/org/springframework/ai/openai/client/AiStreamClient.java rename to spring-ai-core/src/main/java/org/springframework/ai/client/AiStreamClient.java index 6c20260b6..6a655dc92 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/AiStreamClient.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/client/AiStreamClient.java @@ -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 generateStream(String message) { - Prompt prompt = new Prompt(new UserMessage(message)); - return generateStream(prompt); - } - - Flux generateStream(Prompt prompt); + public Flux generateStream(Prompt prompt); } diff --git a/spring-ai-openai/pom.xml b/spring-ai-openai/pom.xml index 500bcd322..6da2bd803 100644 --- a/spring-ai-openai/pom.xml +++ b/spring-ai-openai/pom.xml @@ -27,6 +27,12 @@ ${project.parent.version} + + org.springframework.retry + spring-retry + 2.0.4 + + io.rest-assured json-path @@ -37,12 +43,6 @@ okhttp - - com.theokanning.openai-gpt3-java - service - ${open-ai-client.version} - - com.github.victools jsonschema-generator diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java new file mode 100644 index 000000000..de3460d31 --- /dev/null +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java @@ -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 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 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 messages, + @JsonProperty("model") String model, + @JsonProperty("frequency_penalty") Float frequencyPenalty, + @JsonProperty("logit_bias") Map 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 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 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 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 messages, String model, + List 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 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 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 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 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 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 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 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 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( + @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 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( + @JsonProperty("object") String object, + @JsonProperty("data") List 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 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: + * + * 
{@code List.of("text1", "text2", "text3") or List.of(List.of(1, 2, 3), List.of(3, 4, 5))}
+ */ + public ResponseEntity> embeddings(EmbeddingRequest 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 parseJson(String jsonSchema) { + try { + return new ObjectMapper().readValue(jsonSchema, + new TypeReference>() { + }); + } + catch (Exception e) { + throw new OpenAiApiException("Failed to parse schema: " + jsonSchema, e); + } + } + + private T parseJson(String json, Class type) { + try { + return this.objectMapper.readValue(json, type); + } + catch (Exception e) { + throw new OpenAiApiException("Failed to parse schema: " + json, e); + } + } +} +// @formatter:on diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionResponse.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionResponse.java deleted file mode 100644 index 697363a05..000000000 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionResponse.java +++ /dev/null @@ -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 choices, - - @JsonProperty("created") Integer created, - - @JsonProperty("model") String model, - - @JsonProperty("system_fingerprint") String systemFingerprint, - - @JsonProperty("object") String object, - - @JsonProperty("usage") Map usage) { - public record Choice( - - @JsonProperty("finish_reason") String finishReason, - - @JsonProperty("index") Integer index, - - @JsonProperty("message") OpenAiChatMessage message) { - } -} \ No newline at end of file diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionsRequest.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionsRequest.java deleted file mode 100644 index 102a9598a..000000000 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/ChatCompletionsRequest.java +++ /dev/null @@ -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 messages; - - private final String model; - - private final Integer frequencyPenalty; - - private final Map logitBias; - - private final Integer maxTokens; - - private final Integer n; - - private final Integer presencePenalty; - - private final ResponseFormat responseFormat; - - private final Integer seed; - - private final List stop; - - private final Boolean stream; - - private final Double temperature; - - private final Integer topP; - - private final List 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 getMessages() { - return messages; - } - - public String getModel() { - return model; - } - - public Integer getFrequencyPenalty() { - return frequencyPenalty; - } - - public Map 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 getStop() { - return stop; - } - - public Boolean getStream() { - return stream; - } - - public Double getTemperature() { - return temperature; - } - - public Integer getTopP() { - return topP; - } - - public List getTools() { - return tools; - } - - public String getToolChoice() { - return toolChoice; - } - - public String getUser() { - return user; - } - - public static class Builder { - - @JsonProperty("messages") - private List messages; - - @JsonProperty("model") - private String model; - - @JsonProperty("frequency_penalty") - private Integer frequencyPenalty; - - @JsonProperty("logit_bias") - private Map 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 stop; - - @JsonProperty("stream") - private Boolean stream; - - @JsonProperty("temperature") - private Double temperature; - - @JsonProperty("top_p") - private Integer topP; - - @JsonProperty("tools") - private List tools; - - @JsonProperty("tool_choice") - private String toolChoice; - - @JsonProperty("user") - private String user; - - public Builder messages(List 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 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 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 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 parameters, - - @JsonProperty("arguments") String arguments) { - } - - public record ResponseFormat( - - @JsonProperty("type") String type) { - } - - public record Tool( - - @JsonProperty("function") Function function, - - @JsonProperty("type") String type) { - } - -} diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiChatMessage.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiChatMessage.java deleted file mode 100644 index ee1d14360..000000000 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiChatMessage.java +++ /dev/null @@ -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 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 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 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 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) { - } - -} \ No newline at end of file diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiClient.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiClient.java index e84e674fb..604e20ff2 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiClient.java +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiClient.java @@ -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 messages = prompt.getMessages(); - List theoMessages = messages.stream() - .map(message -> new ChatMessage(message.getMessageTypeValue(), message.getContent())) + List 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 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 chatCompletionChoices = chatCompletionResult.getChoices(); - logger.trace("ChatCompletionChoices: {}", chatCompletionChoices); - List 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 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 getChatCompletionRequest(Prompt prompt) { - - List chatMessages = convertToChatMessages(prompt.getMessages()); - List chatCompletionRequests = new ArrayList<>(); - - ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder() - .model(this.model) - .temperature(this.temperature) - .messages(chatMessages) - .build(); - - chatCompletionRequests.add(chatCompletionRequest); - - return chatCompletionRequests; - } - - private List convertToChatMessages(List messages) { - - List 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 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 generateStream(Prompt prompt) { + List messages = prompt.getMessages(); + + List chatCompletionMessages = messages.stream() + .map(m -> new ChatCompletionMessage(m.getContent(), + ChatCompletionMessage.Role.valueOf(m.getMessageType().getValue()))) + .toList(); + + Flux 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 roleMap = new ConcurrentHashMap<>(); + + // An alternative implementation that returns Flux instead of + // Flux. + // Flux 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 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); + }); + } } diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiSseResponse.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiSseResponse.java deleted file mode 100644 index 78923b61b..000000000 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiSseResponse.java +++ /dev/null @@ -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 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) { - } - } -} \ No newline at end of file diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiStreamClient.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiStreamClient.java deleted file mode 100644 index cd8bf213f..000000000 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/client/OpenAiStreamClient.java +++ /dev/null @@ -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> 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 generateStream(Prompt prompt) { - - List 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)); - } - }); - } - -} diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java index a85fd70cf..8b073324b 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java @@ -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 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 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 embed(String text) { + Assert.notNull(text, "Text must not be null"); + return this.embed(List.of(text)).iterator().next(); + } + + @Override public List> embed(List texts) { - EmbeddingResponse embeddingResponse = embedForResponse(texts); - return embeddingResponse.getData().stream().map(emb -> emb.getEmbedding()).toList(); + Assert.notNull(texts, "Texts must not be null"); + EmbeddingRequest> request = new EmbeddingRequest<>(texts, this.embeddingModelName); + return this.retryTemplate.execute(ctx -> { + EmbeddingList 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 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 data = generateEmbeddingList(nativeEmbeddingResult.getData()); - Map metadata = generateMetadata(nativeEmbeddingResult.getModel(), - nativeEmbeddingResult.getUsage()); - return new EmbeddingResponse(data, metadata); - } + Assert.notNull(texts, "Texts must not be null"); - private List generateEmbeddingList(List nativeData) { - List data = new ArrayList<>(); - for (com.theokanning.openai.embedding.Embedding nativeDatum : nativeData) { - List nativeDatumEmbedding = nativeDatum.getEmbedding(); - int nativeIndex = nativeDatum.getIndex(); - Embedding embedding = new Embedding(nativeDatumEmbedding, nativeIndex); - data.add(embedding); - } - return data; + return this.retryTemplate.execute(ctx -> { + EmbeddingRequest> request = new EmbeddingRequest<>(texts, this.embeddingModelName); + + EmbeddingList 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 metadata = generateMetadata(embeddingResponse.model(), embeddingResponse.usage()); + + List embeddings = embeddingResponse.data() + .stream() + .map(e -> new Embedding(e.embedding(), e.index())) + .toList(); + + return new EmbeddingResponse(embeddings, metadata); + }); } private Map generateMetadata(String model, Usage usage) { Map 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; } diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiGenerationMetadata.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiGenerationMetadata.java index 0b3222de8..271142bc3 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiGenerationMetadata.java +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiGenerationMetadata.java @@ -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; } diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java index 02b02de6e..1b08d449d 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java @@ -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 diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/support/OpenAiResponseHeaderExtractor.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/support/OpenAiResponseHeaderExtractor.java new file mode 100644 index 000000000..be6eb4b4b --- /dev/null +++ b/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/support/OpenAiResponseHeaderExtractor.java @@ -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 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 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 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 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()); + } + + } + + } + +} diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/MockOpenAiTestConfiguration.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/MockOpenAiTestConfiguration.java deleted file mode 100644 index 0822eb3d7..000000000 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/MockOpenAiTestConfiguration.java +++ /dev/null @@ -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. - *

- * 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}. - *

- * 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); - } - -} diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java index 0db246958..1f8767a4d 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java @@ -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); } } diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java index c28a2ef6c..5baf05229 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java @@ -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); diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadataTests.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadata2Tests.java similarity index 50% rename from spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadataTests.java rename to spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadata2Tests.java index 30ce86faf..83011d74b 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadataTests.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientWithGenerationMetadata2Tests.java @@ -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); } } diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java index 368af4a57..2440b998c 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java @@ -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); } diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java index c8685eab6..8f22e2c92 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java @@ -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; diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java index 47e94fcfc..0cddec197 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java @@ -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; } diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 00d839dd6..70fe61656 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -72,6 +72,14 @@ true + + + io.netty + netty-codec-http2 + 4.1.100.Final + + org.springframework.ai @@ -140,12 +148,6 @@ true - - com.squareup.retrofit2 - converter-jackson - 2.9.0 - - diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java index 9b0e1e12c..f2785758f 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java @@ -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()); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiProperties.java index 245c42a12..a86975d39 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiProperties.java @@ -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; } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiPropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiPropertiesTests.java index 74390710a..38d6d3b2a 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiPropertiesTests.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiPropertiesTests.java @@ -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(); diff --git a/spring-ai-vertex-ai/pom.xml b/spring-ai-vertex-ai/pom.xml index 219e74d96..59e687737 100644 --- a/spring-ai-vertex-ai/pom.xml +++ b/spring-ai-vertex-ai/pom.xml @@ -28,22 +28,10 @@ ${project.parent.version} - - com.google.cloud - google-cloud-aiplatform - 3.32.0 - - - commons-logging - commons-logging - - - - org.springframework spring-web - 6.1.1 + ${spring-framework.version} diff --git a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java index 2256fabf3..a8da6bca3 100644 --- a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java +++ b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java index 3bbe8e2ca..bd8add2a1 100644 --- a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java index c6134b479..d2b75b4d3 100644 --- a/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java +++ b/vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-milvus-store/pom.xml b/vector-stores/spring-ai-milvus-store/pom.xml index 24376ae78..a60093dc1 100644 --- a/vector-stores/spring-ai-milvus-store/pom.xml +++ b/vector-stores/spring-ai-milvus-store/pom.xml @@ -57,8 +57,9 @@ org.springframework.ai - spring-ai-openai + + spring-ai-vertex-ai ${parent.version} test diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java index e28ac2136..da97e17f9 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java index 0555b815e..d31b28b31 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java index 58477880e..d092c008b 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java @@ -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"))); } } diff --git a/vector-stores/spring-ai-pinecone/README.md b/vector-stores/spring-ai-pinecone/README.md new file mode 100644 index 000000000..37b1bc022 --- /dev/null +++ b/vector-stores/spring-ai-pinecone/README.md @@ -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 + + spring-snapshots + Spring Snapshots + https://repo.spring.io/snapshot + + false + + +``` + +## Dependencies + +Add these dependencies to your project: + +1. OpenAI: Required for calculating embeddings. + +```xml + + org.springframework.experimental.ai + spring-ai-openai-spring-boot-starter + 0.8.0-SNAPSHOT + +``` + +2. Pinecone + +```xml + + org.springframework.experimental.ai + spring-ai-pinecone + 0.8.0-SNAPSHOT + +``` + +## Sample Code + +To configure Pinecone in your application, you can use the following setup: + +```java +@Bean +public PineconeVectorStoreConfig pineconeVectorStoreConfig() { + + return PineconeVectorStoreConfig.builder() + .withApiKey() + .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 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 results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5)); +``` + +If all goes well, you should retrieve the document containing the text "Spring AI rocks!!". diff --git a/vector-stores/spring-ai-pinecone/pom.xml b/vector-stores/spring-ai-pinecone/pom.xml index 395b7eee0..73d13dab8 100644 --- a/vector-stores/spring-ai-pinecone/pom.xml +++ b/vector-stores/spring-ai-pinecone/pom.xml @@ -36,8 +36,31 @@ io.pinecone pinecone-client ${pinecone.version} + + + io.netty + netty-codec + + + io.netty + netty-codec-http2 + + + + + io.netty + netty-codec + 4.1.101.Final + + + io.netty + netty-codec-http2 + 4.1.100.Final + + + com.google.protobuf protobuf-java-util @@ -45,19 +68,19 @@ - - org.springframework.ai - transformers-embedding - ${parent.version} - test - + + org.springframework.ai + transformers-embedding + ${parent.version} + test + - - org.springframework.ai - spring-ai-test - ${parent.version} - test - + + org.springframework.ai + spring-ai-test + ${parent.version} + test + org.springframework.boot