Make sure that managed QuarkusAiServiceContext is always used

Author: geoand

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/AiServicesProcessor.java

@@ -115,8 +115,10 @@
 import io.quarkiverse.langchain4j.spi.PromptTemplateFactoryContentFilterProvider;
 import io.quarkus.arc.Arc;
 import io.quarkus.arc.ArcContainer;
+import io.quarkus.arc.DefaultBean;
 import io.quarkus.arc.InstanceHandle;
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
+import io.quarkus.arc.deployment.AnnotationsTransformerBuildItem;
 import io.quarkus.arc.deployment.CustomScopeAnnotationsBuildItem;
 import io.quarkus.arc.deployment.GeneratedBeanBuildItem;
 import io.quarkus.arc.deployment.GeneratedBeanGizmoAdaptor;
@@ -324,6 +326,7 @@ private static Set<String> gatherToolNames(ClassInfo toolClass) {
     @BuildStep
     public void findDeclarativeServices(CombinedIndexBuildItem indexBuildItem,
             CustomScopeAnnotationsBuildItem customScopes,
+            List<AnnotationsImpliesAiServiceBuildItem> annotationsImpliesAiServiceItems,
             BuildProducer<RequestChatModelBeanBuildItem> requestChatModelBeanProducer,
             BuildProducer<RequestModerationModelBeanBuildItem> requestModerationModelBeanProducer,
             BuildProducer<RequestImageModelBeanBuildItem> requestImageModelBeanProducer,
@@ -338,11 +341,29 @@ public void findDeclarativeServices(CombinedIndexBuildItem indexBuildItem,
         Set<String> imageModelNames = new HashSet<>();
         List<ToolProviderInfo> toolProviderInfos = new ArrayList<>();
         ClassOutput generatedClassOutput = new GeneratedClassGizmoAdaptor(generatedClassProducer, true);
-        for (AnnotationInstance instance : index.getAnnotations(LangChain4jDotNames.REGISTER_AI_SERVICES)) {
+
+        Set<DotName> annotationsThatImplyAiService = annotationsImpliesAiServiceItems.stream().flatMap(
+                bi -> bi.getAnnotationNames().stream())
+                .collect(Collectors.toSet());
+
+        Collection<AnnotationInstance> registerAiServicesInstances = new ArrayList<>(
+                index.getAnnotations(REGISTER_AI_SERVICES));
+
+        Set<AnnotationInstance> impliedRegisterAiServiceInstance = determinedImpliedRegisterAiService(
+                annotationsThatImplyAiService, index);
+        Set<DotName> impliedRegisterAiServiceTarget = impliedRegisterAiServiceInstance.stream()
+                .map(ai -> ai.target().asClass().name()).collect(Collectors.toSet());
+        registerAiServicesInstances.addAll(impliedRegisterAiServiceInstance);
+        Set<DotName> alreadyHandled = new HashSet<>();
+        for (AnnotationInstance instance : registerAiServicesInstances) {
             if (instance.target().kind() != AnnotationTarget.Kind.CLASS) {
                 continue; // should never happen
             }
             ClassInfo declarativeAiServiceClassInfo = instance.target().asClass();
+            if (alreadyHandled.contains(declarativeAiServiceClassInfo.name())) {
+                continue;
+            }
+            alreadyHandled.add(declarativeAiServiceClassInfo.name());
 
             DotName chatLanguageModelSupplierClassDotName = getSupplierDotName(instance.value("chatLanguageModelSupplier"),
                     LangChain4jDotNames.BEAN_CHAT_MODEL_SUPPLIER,
@@ -462,7 +483,9 @@ public void findDeclarativeServices(CombinedIndexBuildItem indexBuildItem,
                             toolHallucinationStrategy(instance),
                             classInputGuardrails(declarativeAiServiceClassInfo, index),
                             classOutputGuardrails(declarativeAiServiceClassInfo, index),
-                            maxSequentialToolInvocations));
+                            maxSequentialToolInvocations,
+                            // we need to make these @DefaultBean because there could be other CDI beans of the same type that need to take precedence
+                            impliedRegisterAiServiceTarget.contains(declarativeAiServiceClassInfo.name())));
 
         }
         toolProviderProducer.produce(new ToolProviderMetaBuildItem(toolProviderInfos));
@@ -480,6 +503,35 @@ public void findDeclarativeServices(CombinedIndexBuildItem indexBuildItem,
         }
     }
 
+    private static Set<AnnotationInstance> determinedImpliedRegisterAiService(Set<DotName> annotationsThatImplyAiService,
+            IndexView index) {
+        Set<AnnotationInstance> impliedDefaultRegisterAiService = new HashSet<>();
+        for (DotName ann : annotationsThatImplyAiService) {
+            index.getAnnotations(ann).forEach(instance -> {
+                ClassInfo ci;
+                switch (instance.target().kind()) {
+                    case METHOD -> {
+                        ci = instance.target().asMethod().declaringClass();
+                    }
+                    case CLASS -> {
+                        ci = instance.target().asClass();
+                    }
+                    case FIELD -> {
+                        ci = instance.target().asField().declaringClass();
+                    }
+                    default -> {
+                        ci = null;
+                    }
+                }
+                if (ci == null) {
+                    return;
+                }
+                impliedDefaultRegisterAiService.add(AnnotationInstance.builder(REGISTER_AI_SERVICES).buildWithTarget(ci));
+            });
+        }
+        return impliedDefaultRegisterAiService;
+    }
+
     private static String chatModelName(AnnotationInstance instance, DotName chatLanguageModelSupplierClassDotName,
             DotName streamingChatLanguageModelSupplierClassDotName, Set<String> chatModelNames) {
         String chatModelName = NamedConfigUtil.DEFAULT_NAME;
@@ -793,6 +845,7 @@ public void handleDeclarativeServices(AiServicesRecorder recorder,
             String moderationModelName = bi.getModerationModelName();
             SyntheticBeanBuildItem.ExtendedBeanConfigurator configurator = SyntheticBeanBuildItem
                     .configure(QuarkusAiServiceContext.class)
+                    .unremovable()
                     .forceApplicationClass()
                     .createWith(recorder.createDeclarativeAiService(
                             new DeclarativeAiServiceCreateInfo(
@@ -1166,6 +1219,11 @@ AnnotationsImpliesAiServiceBuildItem implyAiService() {
                         LangChain4jDotNames.MODERATE));
     }
 
+    @BuildStep
+    public void annotationTransformations(BuildProducer<AnnotationsTransformerBuildItem> producer) {
+
+    }
+
     @BuildStep
     @Record(ExecutionTime.STATIC_INIT)
     public void handleAiServices(
@@ -1319,6 +1377,9 @@ public void handleAiServices(
                             classCreator.addAnnotation(
                                     AnnotationInstance.builder(NAMED).add("value", matchingBI.getBeanName().get()).build());
                         }
+                        if (matchingBI.isMakeDefaultBean()) {
+                            classCreator.addAnnotation(DefaultBean.class);
+                        }
                     }
 
                     FieldDescriptor contextField = classCreator.getFieldCreator("context", QuarkusAiServiceContext.class)

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/DeclarativeAiServiceBuildItem.java

@@ -34,6 +34,7 @@ public final class DeclarativeAiServiceBuildItem extends MultiBuildItem {
     private final DeclarativeAiServiceInputGuardrails inputGuardrails;
     private final DeclarativeAiServiceOutputGuardrails outputGuardrails;
     private final Integer maxSequentialToolInvocations;
+    private final boolean makeDefaultBean;
 
     public DeclarativeAiServiceBuildItem(
             ClassInfo serviceClassInfo,
@@ -55,7 +56,7 @@ public DeclarativeAiServiceBuildItem(
             DotName toolHallucinationStrategyClassDotName,
             DeclarativeAiServiceInputGuardrails inputGuardrails,
             DeclarativeAiServiceOutputGuardrails outputGuardrails,
-            Integer maxSequentialToolInvocations) {
+            Integer maxSequentialToolInvocations, boolean makeDefaultBean) {
         this.serviceClassInfo = serviceClassInfo;
         this.chatLanguageModelSupplierClassDotName = chatLanguageModelSupplierClassDotName;
         this.streamingChatLanguageModelSupplierClassDotName = streamingChatLanguageModelSupplierClassDotName;
@@ -76,6 +77,7 @@ public DeclarativeAiServiceBuildItem(
         this.inputGuardrails = inputGuardrails;
         this.outputGuardrails = outputGuardrails;
         this.maxSequentialToolInvocations = maxSequentialToolInvocations;
+        this.makeDefaultBean = makeDefaultBean;
     }
 
     public ClassInfo getServiceClassInfo() {
@@ -154,6 +156,10 @@ public DeclarativeAiServiceOutputGuardrails getOutputGuardrails() {
         return outputGuardrails;
     }
 
+    public boolean isMakeDefaultBean() {
+        return makeDefaultBean;
+    }
+
     public record DeclarativeAiServiceInputGuardrails(List<ClassInfo> inputGuardrailClassInfos) {
         public List<String> asClassNames() {
             return this.inputGuardrailClassInfos.stream()

core/runtime/src/main/java/io/quarkiverse/langchain4j/QuarkusAiServiceContextFactory.java

@@ -3,11 +3,20 @@
 import dev.langchain4j.service.AiServiceContext;
 import dev.langchain4j.spi.services.AiServiceContextFactory;
 import io.quarkiverse.langchain4j.runtime.aiservice.QuarkusAiServiceContext;
+import io.quarkiverse.langchain4j.runtime.aiservice.QuarkusAiServiceContextQualifier;
+import io.quarkus.arc.Arc;
+import io.quarkus.arc.InstanceHandle;
 
 public class QuarkusAiServiceContextFactory implements AiServiceContextFactory {
 
     @Override
     public AiServiceContext create(Class<?> aiServiceClass) {
-        return new QuarkusAiServiceContext(aiServiceClass);
+        InstanceHandle<QuarkusAiServiceContext> instance = Arc.container().instance(QuarkusAiServiceContext.class,
+                QuarkusAiServiceContextQualifier.Literal.of(
+                        aiServiceClass.getName()));
+        if (instance.isAvailable()) {
+            return instance.get();
+        }
+        return new QuarkusAiServiceContext(aiServiceClass); // just create a default context
     }
 }

model-providers/openai/openai-vanilla/deployment/src/test/java/org/acme/examples/aiservices/AiServicesTest.java

@@ -58,6 +58,7 @@
 import io.opentelemetry.instrumentation.annotations.SpanAttribute;
 import io.quarkiverse.langchain4j.openai.testing.internal.OpenAiBaseTest;
 import io.quarkiverse.langchain4j.runtime.LangChain4jUtil;
+import io.quarkus.arc.Arc;
 import io.quarkus.test.QuarkusUnitTest;
 
 public class AiServicesTest extends OpenAiBaseTest {
@@ -68,7 +69,8 @@ public class AiServicesTest extends OpenAiBaseTest {
                     () -> ShrinkWrap.create(JavaArchive.class)
                             .addAsResource("messages/recipe-user.txt")
                             .addAsResource("messages/translate-user.txt")
-                            .addAsResource("messages/translate-system"));
+                            .addAsResource("messages/translate-system"))
+            .overrideRuntimeConfigKey("quarkus.langchain4j.openai.api-key", "whatever");
 
     private OpenAiChatModel createChatModel() {
         return OpenAiChatModel.builder().baseUrl(resolvedWiremockUrl("/v1"))
@@ -93,6 +95,11 @@ private static MessageWindowChatMemory createChatMemory() {
     void setup() {
         resetRequests();
         resetMappings();
+        clearDefaultMemory();
+    }
+
+    private void clearDefaultMemory() {
+        Arc.container().instance(ChatMemoryStore.class).get().deleteMessages("default");
     }
 
     interface Assistant {
@@ -312,6 +319,7 @@ void test_create_recipe_from_list_of_ingredients() throws IOException {
 
                     resetRequests();
                     resetMappings();
+                    clearDefaultMemory();
 
                     setChatCompletionMessageContent(value);
                     Chef chef = AiServices.create(Chef.class, createChatModel());