diff --git a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/HuggingfaceTestConfiguration.java b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/HuggingfaceTestConfiguration.java index 011650141..8e2a90d9d 100644 --- a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/HuggingfaceTestConfiguration.java +++ b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/HuggingfaceTestConfiguration.java @@ -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; } diff --git a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java index dfb46b4a0..48f32b6ba 100644 --- a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java +++ b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java @@ -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); } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/huggingface/HuggingfaceChatAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/huggingface/HuggingfaceChatAutoConfigurationIT.java new file mode 100644 index 000000000..300962345 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/huggingface/HuggingfaceChatAutoConfigurationIT.java @@ -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 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); + }); + } + +}