AiService support @Profile (#41)

Closes langchain4j/langchain4j#1430 Support `@AiService` with `@Profile`.
langchain4j · Sep 18, 2024 · 3648f2c · 3648f2c
1 parent dfd6c11
commit 3648f2c
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 19 deletions.
diff --git a/langchain4j-spring-boot-starter/pom.xml b/langchain4j-spring-boot-starter/pom.xml
@@ -69,19 +69,6 @@
             <scope>test</scope>
         </dependency>
 
-        <dependency>
-            <groupId>org.tinylog</groupId>
-            <artifactId>tinylog-impl</artifactId>
-            <version>2.6.2</version>
-            <scope>test</scope>
-        </dependency>
-        <dependency>
-            <groupId>org.tinylog</groupId>
-            <artifactId>slf4j-tinylog</artifactId>
-            <version>2.6.2</version>
-            <scope>test</scope>
-        </dependency>
-
     </dependencies>
 
     <licenses>

diff --git a/...-boot-starter/src/main/java/dev/langchain4j/service/spring/AiServiceScannerProcessor.java b/...-boot-starter/src/main/java/dev/langchain4j/service/spring/AiServiceScannerProcessor.java
@@ -6,18 +6,20 @@
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigurationPackages;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.context.EnvironmentAware;
 import org.springframework.context.annotation.ComponentScan;
+import org.springframework.context.annotation.Profile;
 import org.springframework.core.annotation.AnnotationUtils;
+import org.springframework.core.env.Environment;
 import org.springframework.stereotype.Component;
 import org.springframework.util.ClassUtils;
 
-import java.util.Collections;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 
 @Component
-public class AiServiceScannerProcessor implements BeanDefinitionRegistryPostProcessor {
+public class AiServiceScannerProcessor implements BeanDefinitionRegistryPostProcessor, EnvironmentAware {
+
+    private Environment environment;
 
     @Override
     public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException {
@@ -26,6 +28,8 @@ public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) t
         for (String basePackage : basePackages) {
             classPathAiServiceScanner.scan(basePackage);
         }
+
+        filterBeanDefinitions(registry);
     }
 
     private Set<String> getBasePackages(ConfigurableListableBeanFactory beanFactory) {
@@ -63,4 +67,28 @@ private Set<String> getBasePackages(ConfigurableListableBeanFactory beanFactory)
 
         return basePackages;
     }
+
+    private void filterBeanDefinitions(BeanDefinitionRegistry registry) {
+        Arrays.stream(registry.getBeanDefinitionNames())
+                .filter(beanName -> {
+                    try {
+                        Class<?> beanClass = Class.forName(registry.getBeanDefinition(beanName).getBeanClassName());
+                        if (beanClass.isAnnotationPresent(AiService.class) && beanClass.isAnnotationPresent(Profile.class)) {
+                            Profile profileAnnotation = beanClass.getAnnotation(Profile.class);
+                            String[] profiles = profileAnnotation.value();
+
+                            return !environment.matchesProfiles(profiles);
+                        } else {
+                            return false;
+                        }
+                    } catch (ClassNotFoundException e) {
+                        return false;
+                    }
+                }).forEach(registry::removeBeanDefinition);
+    }
+
+    @Override
+    public void setEnvironment(Environment environment) {
+        this.environment = environment;
+    }
 }
diff --git a/...pring-boot-starter/src/main/java/dev/langchain4j/service/spring/AiServicesAutoConfig.java b/...pring-boot-starter/src/main/java/dev/langchain4j/service/spring/AiServicesAutoConfig.java
@@ -17,7 +17,10 @@
 import org.springframework.context.annotation.Bean;
 
 import java.lang.reflect.Method;
-import java.util.*;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
 
 import static dev.langchain4j.exception.IllegalConfigurationException.illegalConfiguration;
 import static dev.langchain4j.internal.Exceptions.illegalArgument;

diff --git a/...ava/dev/langchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfiles.java b/...ava/dev/langchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfiles.java
@@ -0,0 +1,11 @@
+package dev.langchain4j.service.spring.mode.automatic.withProfiles;
+
+import dev.langchain4j.service.spring.AiService;
+import org.springframework.context.annotation.Profile;
+
+@AiService
+@Profile("!test")
+public interface AiServiceWithProfiles {
+
+    String chat(String userMessage);
+}
diff --git a/...gchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfilesApplication.java b/...gchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfilesApplication.java
@@ -0,0 +1,20 @@
+package dev.langchain4j.service.spring.mode.automatic.withProfiles;
+
+import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.openai.OpenAiChatModel;
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.context.annotation.Bean;
+
+@SpringBootApplication
+public class AiServiceWithProfilesApplication {
+
+    @Bean
+    ChatLanguageModel chatLanguageModel() {
+        return OpenAiChatModel.withApiKey(System.getenv("OPENAI_API_KEY"));
+    }
+
+    public static void main(String[] args) {
+        SpringApplication.run(AiServiceWithProfilesApplication.class, args);
+    }
+}
diff --git a/...a/dev/langchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfilesIT.java b/...a/dev/langchain4j/service/spring/mode/automatic/withProfiles/AiServiceWithProfilesIT.java
@@ -0,0 +1,37 @@
+package dev.langchain4j.service.spring.mode.automatic.withProfiles;
+
+import dev.langchain4j.service.spring.AiServicesAutoConfig;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.BeansException;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class AiServiceWithProfilesIT {
+
+    ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+            .withConfiguration(AutoConfigurations.of(AiServicesAutoConfig.class));
+
+    @Test
+    void should_not_create_ai_service() {
+        contextRunner
+                .withPropertyValues(
+                        "spring.profiles.active=test"
+                )
+                .withUserConfiguration(AiServiceWithProfilesApplication.class)
+                .run(context -> assertThatThrownBy(() -> context.getBean(AiServiceWithProfiles.class))
+                        .isInstanceOf(BeansException.class));
+    }
+
+    @Test
+    void should_create_ai_service() {
+        contextRunner
+                .withPropertyValues(
+                        "spring.profiles.active=dev"
+                )
+                .withUserConfiguration(AiServiceWithProfilesApplication.class)
+                .run(context -> assertThat(context.getBean(AiServiceWithProfiles.class)).isNotNull());
+    }
+}