add OpenAiApi tests

This commit is contained in:
Christian Tzolov
2024-02-01 19:13:57 +01:00
parent b49458b3ee
commit d71f90fd56
3 changed files with 160 additions and 1 deletions

View File

@@ -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());
}
/**

View File

@@ -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<ChatCompletion> 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<ChatCompletionChunk> 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<EmbeddingList<Embedding>> response = openAiApi
.embeddings(new OpenAiApi.EmbeddingRequest<String>("Hello world"));
assertThat(response).isNotNull();
assertThat(response.getBody().data()).hasSize(1);
assertThat(response.getBody().data().get(0).embedding()).hasSize(1536);
}
}

View File

@@ -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();
}
}