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

Improve Hugging Face ITs

Browse Source
This commit is contained in:
Christian Tzolov
2024-06-13 22:13:40 +02:00
parent 2a592d4e84
commit be9e4a828a
3 changed files with 94 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/HuggingfaceTestConfiguration.java
View File

@@ -29,10 +29,9 @@ public class HuggingfaceTestConfiguration {
throw new IllegalArgumentException(
"You must provide an API key. Put it in an environment variable under the name HUGGINGFACE_API_KEY");
}
// Created aws-mistral-7b-instruct-v0-1-805 via
// https://ui.endpoints.huggingface.co/
// Created aws-mistral-7b-instruct and update the HUGGINGFACE_CHAT_URL
HuggingfaceChatModel huggingfaceChatModel = new HuggingfaceChatModel(apiKey,
"https://f6hg7b3cvlmntp5i.us-east-1.aws.endpoints.huggingface.cloud");
System.getenv("HUGGINGFACE_CHAT_URL"));
return huggingfaceChatModel;
}

22
models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java
View File

@@ -15,21 +15,19 @@
*/
package org.springframework.ai.huggingface.client;
import org.junit.jupiter.api.Disabled;
import static org.assertj.core.api.Assertions.assertThat;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.huggingface.HuggingfaceChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.huggingface.HuggingfaceChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import static org.assertj.core.api.Assertions.assertThat;
@Disabled("Until a valid inference endpoint is available for the provided HUGGINGFACE_API_KEY ")
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "HUGGINGFACE_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "HUGGINGFACE_CHAT_URL", matches = ".+")
public class ClientIT {
@Autowired
@@ -49,16 +47,14 @@ public class ClientIT {
ChatResponse chatResponse = huggingfaceChatModel.call(prompt);
assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty();
String expectedResponse = """
```json
{
"name": "John",
"lastname": "Smith",
"address": "#1 Samuel St."
}
```""";
"name": "John",
"lastname": "Smith",
"address": "#1 Samuel St."
}""";
assertThat(chatResponse.getResult().getOutput().getContent()).isEqualTo(expectedResponse);
assertThat(chatResponse.getResult().getOutput().getMetadata()).containsKey("generated_tokens");
assertThat(chatResponse.getResult().getOutput().getMetadata()).containsEntry("generated_tokens", 39);
assertThat(chatResponse.getResult().getOutput().getMetadata()).containsEntry("generated_tokens", 32);
}

83
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/huggingface/HuggingfaceChatAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,83 @@
/*
* 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.autoconfigure.huggingface;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.huggingface.HuggingfaceChatModel;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import reactor.core.publisher.Flux;
@EnabledIfEnvironmentVariable(named = "HUGGINGFACE_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "HUGGINGFACE_CHAT_URL", matches = ".+")
public class HuggingfaceChatAutoConfigurationIT {
private static final Log logger = LogFactory.getLog(HuggingfaceChatAutoConfigurationIT.class);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner().withPropertyValues(
// @formatter:off
"spring.ai.huggingface.chat.api-key=" + System.getenv("HUGGINGFACE_API_KEY"),
"spring.ai.huggingface.chat.url=" + System.getenv("HUGGINGFACE_CHAT_URL"))
// @formatter:on
.withConfiguration(AutoConfigurations.of(HuggingfaceChatAutoConfiguration.class));
@Test
void generate() {
contextRunner.run(context -> {
HuggingfaceChatModel chatModel = context.getBean(HuggingfaceChatModel.class);
String response = chatModel.call("Hello");
assertThat(response).isNotEmpty();
logger.info("Response: " + response);
});
}
@Disabled("Until streaming support is added")
@Test
void generateStreaming() {
contextRunner.run(context -> {
HuggingfaceChatModel chatModel = context.getBean(HuggingfaceChatModel.class);
Flux<ChatResponse> responseFlux = chatModel.stream(new Prompt(new UserMessage("Hello")));
String response = responseFlux.collectList()
.block()
.stream()
.map(ChatResponse::getResults)
.flatMap(List::stream)
.map(Generation::getOutput)
.map(AssistantMessage::getContent)
.collect(Collectors.joining());
assertThat(response).isNotEmpty();
logger.info("Response: " + response);
});
}
}