From d71f90fd563254cc32447bf2a6c64d4b6b293472 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 1 Feb 2024 19:13:57 +0100 Subject: [PATCH] add OpenAiApi tests --- .../ai/openai/api/OpenAiApi.java | 12 ++- .../ai/openai/chat/api/OpenAiApiIT.java | 75 +++++++++++++++++++ .../chat/api/RestClientBuilderTests.java | 74 ++++++++++++++++++ 3 files changed, 160 insertions(+), 1 deletion(-) create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/OpenAiApiIT.java create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/RestClientBuilderTests.java diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java index 6e5632f78..097234a04 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java @@ -66,7 +66,17 @@ public class OpenAiApi { * @param openAiToken OpenAI apiKey. */ public OpenAiApi(String openAiToken) { - this(DEFAULT_BASE_URL, openAiToken, RestClient.builder()); + this(DEFAULT_BASE_URL, openAiToken); + } + + /** + * Create a new chat completion api. + * + * @param baseUrl api base URL. + * @param openAiToken OpenAI apiKey. + */ + public OpenAiApi(String baseUrl, String openAiToken) { + this(baseUrl, openAiToken, RestClient.builder()); } /** diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/OpenAiApiIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/OpenAiApiIT.java new file mode 100644 index 000000000..af0acbd97 --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/OpenAiApiIT.java @@ -0,0 +1,75 @@ +/* + * Copyright 2024-2024 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.chat.api; + +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import reactor.core.publisher.Flux; + +import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionChunk; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest; +import org.springframework.ai.openai.api.OpenAiApi.Embedding; +import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList; +import org.springframework.http.ResponseEntity; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +public class OpenAiApiIT { + + OpenAiApi openAiApi = new OpenAiApi(System.getenv("OPENAI_API_KEY")); + + @Test + void chatCompletionEntity() { + ChatCompletionMessage chatCompletionMessage = new ChatCompletionMessage("Hello world", Role.USER); + ResponseEntity response = openAiApi.chatCompletionEntity( + new ChatCompletionRequest(List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8f, false)); + + assertThat(response).isNotNull(); + assertThat(response.getBody()).isNotNull(); + } + + @Test + void chatCompletionStream() { + ChatCompletionMessage chatCompletionMessage = new ChatCompletionMessage("Hello world", Role.USER); + Flux response = openAiApi.chatCompletionStream( + new ChatCompletionRequest(List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8f, true)); + + assertThat(response).isNotNull(); + assertThat(response.collectList().block()).isNotNull(); + } + + @Test + void embeddings() { + ResponseEntity> response = openAiApi + .embeddings(new OpenAiApi.EmbeddingRequest("Hello world")); + + assertThat(response).isNotNull(); + assertThat(response.getBody().data()).hasSize(1); + assertThat(response.getBody().data().get(0).embedding()).hasSize(1536); + } + +} diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/RestClientBuilderTests.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/RestClientBuilderTests.java new file mode 100644 index 000000000..f19ce9c5c --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/RestClientBuilderTests.java @@ -0,0 +1,74 @@ +/* + * Copyright 2024-2024 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.chat.api; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; +import org.springframework.http.client.SimpleClientHttpRequestFactory; +import org.springframework.web.client.RestClient; +import org.springframework.web.client.RestClient.Builder; +import org.springframework.web.client.RestTemplate; +import org.springframework.web.util.DefaultUriBuilderFactory; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +public class RestClientBuilderTests { + + public static final String BASE_URL = "https://dog.ceo"; + + @Test + public void test1() { + test(RestClient.builder(), BASE_URL); + } + + @Test + @Disabled("RestClient.builder(restTemplate) bug: https://github.com/spring-projects/spring-framework/issues/32180") + public void test2() { + RestTemplate restTemplate = new RestTemplate(); + test(RestClient.builder(restTemplate), BASE_URL); + } + + @Test + public void test3() { + RestTemplate restTemplate = new RestTemplate(); + restTemplate.setUriTemplateHandler(new DefaultUriBuilderFactory(BASE_URL)); + test(RestClient.builder(restTemplate), BASE_URL); + } + + @Test + public void test4() { + var clientHttpRequestFactory = new SimpleClientHttpRequestFactory(); + clientHttpRequestFactory.setConnectTimeout(5000); + // clientHttpRequestFactory.setProxy(new Proxy(Type.HTTP, + // InetSocketAddress.createUnresolved("localhost", 80))); + RestClient.Builder builder = RestClient.builder().requestFactory(clientHttpRequestFactory); + test(builder, BASE_URL); + } + + private void test(Builder restClientBuilder, String baseUrl) { + var restClient = restClientBuilder.baseUrl(baseUrl).build(); + String res = restClient.get().uri("/api/breeds/list/all").retrieve().body(String.class); + + assertThat(res).isNotNull(); + } + +}