AiService support @Profile (#41)

langchain4j · Sep 18, 2024 · cb78ad3 · cb78ad3
1 parent 3648f2c
commit cb78ad3
Showing 1 changed file with 15 additions and 11 deletions.
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
@@ -1,6 +1,7 @@
 package dev.langchain4j.service.spring;
 
 import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
@@ -29,7 +30,7 @@ public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) t
             classPathAiServiceScanner.scan(basePackage);
         }
 
-        filterBeanDefinitions(registry);
+        removeAiServicesWithInactiveProfiles(registry);
     }
 
     private Set<String> getBasePackages(ConfigurableListableBeanFactory beanFactory) {
@@ -68,22 +69,25 @@ private Set<String> getBasePackages(ConfigurableListableBeanFactory beanFactory)
         return basePackages;
     }
 
-    private void filterBeanDefinitions(BeanDefinitionRegistry registry) {
+    private void removeAiServicesWithInactiveProfiles(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;
+                        BeanDefinition beanDefinition = registry.getBeanDefinition(beanName);
+                        if (beanDefinition.getBeanClassName() != null) {
+                            Class<?> beanClass = Class.forName(beanDefinition.getBeanClassName());
+                            if (beanClass.isAnnotationPresent(AiService.class)
+                                    && beanClass.isAnnotationPresent(Profile.class)) {
+                                Profile profileAnnotation = beanClass.getAnnotation(Profile.class);
+                                String[] profiles = profileAnnotation.value();
+                                return !environment.matchesProfiles(profiles);
+                            }
                         }
                     } catch (ClassNotFoundException e) {
-                        return false;
+                        // should not happen
                     }
+
+                    return false;
                 }).forEach(registry::removeBeanDefinition);
     }