From 66b57cd4feba406e4b8f6821c88d914ac5a5cec1 Mon Sep 17 00:00:00 2001
From: Lester Solbakken <lester.solbakken@gmail.com>
Date: Tue, 16 Apr 2024 13:45:04 +0200
Subject: Use model-integration as provided and specify exported classes in
 PlatformBundles

---
 .../java/com/yahoo/vespa/model/container/PlatformBundles.java | 11 ++++++++++-
 .../xml/BundleInstantiationSpecificationBuilder.java          |  9 +++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

(limited to 'config-model')

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/PlatformBundles.java b/config-model/src/main/java/com/yahoo/vespa/model/container/PlatformBundles.java
index e801884a73a..468cf8dd961 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/PlatformBundles.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/PlatformBundles.java
@@ -38,6 +38,7 @@ public class PlatformBundles {
 
     public static final Path LIBRARY_PATH = Paths.get(Defaults.getDefaults().underVespaHome("lib/jars"));
     public static final String SEARCH_AND_DOCPROC_BUNDLE = BundleInstantiationSpecification.CONTAINER_SEARCH_AND_DOCPROC;
+    public static final String MODEL_INTEGRATION_BUNDLE = BundleInstantiationSpecification.MODEL_INTEGRATION;
 
     // Bundles that must be loaded for all container types.
     public static final Set<Path> COMMON_VESPA_BUNDLES = toBundlePaths(
@@ -88,6 +89,10 @@ public class PlatformBundles {
         return searchAndDocprocComponents.contains(className);
     }
 
+    public static boolean isModelIntegrationClass(String className) {
+        return modelIntegrationComponents.contains(className);
+    }
+
     // This is a hack to allow users to declare components from the search-and-docproc bundle without naming the bundle.
     private static final Set<String> searchAndDocprocComponents = Set.of(
             com.yahoo.docproc.AbstractConcreteDocumentFactory.class.getName(),
@@ -149,7 +154,11 @@ public class PlatformBundles {
             com.yahoo.vespa.streamingvisitors.MetricsSearcher.class.getName(),
             com.yahoo.vespa.streamingvisitors.StreamingBackend.class.getName(),
             ai.vespa.search.llm.LLMSearcher.class.getName(),
-            ai.vespa.search.llm.RAGSearcher.class.getName(),
+            ai.vespa.search.llm.RAGSearcher.class.getName()
+    );
+
+    // This is a hack to allow users to declare components from the model-integration bundle without naming the bundle.
+    private static final Set<String> modelIntegrationComponents = Set.of(
             ai.vespa.llm.clients.OpenAI.class.getName(),
             ai.vespa.llm.clients.LocalLLM.class.getName()
     );
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/xml/BundleInstantiationSpecificationBuilder.java b/config-model/src/main/java/com/yahoo/vespa/model/container/xml/BundleInstantiationSpecificationBuilder.java
index 7e14eafc2ee..1323506eaeb 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/xml/BundleInstantiationSpecificationBuilder.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/xml/BundleInstantiationSpecificationBuilder.java
@@ -26,17 +26,18 @@ public class BundleInstantiationSpecificationBuilder {
         BundleInstantiationSpecification instSpec = new BundleInstantiationSpecification(id, classId, bundle);
         validate(instSpec);
 
-        return bundle == null ? setBundleForSearchAndDocprocComponents(instSpec) : instSpec;
+        return bundle == null ? setBundleForComponent(instSpec) : instSpec;
     }
 
-    private static BundleInstantiationSpecification setBundleForSearchAndDocprocComponents(BundleInstantiationSpecification spec) {
+    private static BundleInstantiationSpecification setBundleForComponent(BundleInstantiationSpecification spec) {
         if (PlatformBundles.isSearchAndDocprocClass(spec.getClassName()))
             return spec.inBundle(PlatformBundles.SEARCH_AND_DOCPROC_BUNDLE);
+        else if (PlatformBundles.isModelIntegrationClass(spec.getClassName()))
+            return spec.inBundle(PlatformBundles.MODEL_INTEGRATION_BUNDLE);
         else
             return spec;
     }
 
-
     private static void validate(BundleInstantiationSpecification instSpec) {
         List<String> forbiddenClasses = List.of(SearchHandler.HANDLER_CLASSNAME, PROCESSING_HANDLER_CLASS);
 
@@ -47,7 +48,7 @@ public class BundleInstantiationSpecificationBuilder {
         }
     }
 
-    //null if missing
+    // null if missing
     private static ComponentSpecification getComponentSpecification(Element spec, String attributeName) {
         return (spec.hasAttribute(attributeName)) ?
                 new ComponentSpecification(spec.getAttribute(attributeName)) :
-- 
cgit v1.2.3