Merge pull request #30773 from vespa-engine/lesters/update-platform-bundles-for-rag-2

Lesters/update platform bundles for rag 2

16 files changed, 175 insertions, 53 deletions

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 3252077a79d..baa9c4b9a48 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
@@ -145,7 +145,10 @@ public class PlatformBundles {
             com.yahoo.search.searchers.CacheControlSearcher.class.getName(),
             com.yahoo.search.searchers.RateLimitingSearcher.class.getName(),
             com.yahoo.vespa.streamingvisitors.MetricsSearcher.class.getName(),
-            com.yahoo.vespa.streamingvisitors.StreamingBackend.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.llm.clients.OpenAI.class.getName()
     );
 
 }
diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 257dd364000..e74fe22c588 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -9160,7 +9160,7 @@
     ],
     "methods" : [
       "public void <init>()",
-      "public void <init>(ai.vespa.llm.LlmClientConfig, com.yahoo.container.jdisc.secretstore.SecretStore)",
+      "public void <init>(ai.vespa.llm.clients.LlmClientConfig, com.yahoo.container.jdisc.secretstore.SecretStore)",
       "protected java.lang.String getApiKey(ai.vespa.llm.InferenceParameters)",
       "protected void setApiKey(ai.vespa.llm.InferenceParameters)",
       "protected java.lang.String getEndpoint()",
@@ -9168,6 +9168,67 @@
     ],
     "fields" : [ ]
   },
+  "ai.vespa.llm.clients.LlmClientConfig$Builder" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "com.yahoo.config.ConfigInstance$Builder"
+    ],
+    "attributes" : [
+      "public",
+      "final"
+    ],
+    "methods" : [
+      "public void <init>()",
+      "public void <init>(ai.vespa.llm.clients.LlmClientConfig)",
+      "public ai.vespa.llm.clients.LlmClientConfig$Builder apiKeySecretName(java.lang.String)",
+      "public ai.vespa.llm.clients.LlmClientConfig$Builder endpoint(java.lang.String)",
+      "public final boolean dispatchGetConfig(com.yahoo.config.ConfigInstance$Producer)",
+      "public final java.lang.String getDefMd5()",
+      "public final java.lang.String getDefName()",
+      "public final java.lang.String getDefNamespace()",
+      "public final boolean getApplyOnRestart()",
+      "public final void setApplyOnRestart(boolean)",
+      "public ai.vespa.llm.clients.LlmClientConfig build()"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.clients.LlmClientConfig$Producer" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "com.yahoo.config.ConfigInstance$Producer"
+    ],
+    "attributes" : [
+      "public",
+      "interface",
+      "abstract"
+    ],
+    "methods" : [
+      "public abstract void getConfig(ai.vespa.llm.clients.LlmClientConfig$Builder)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.llm.clients.LlmClientConfig" : {
+    "superClass" : "com.yahoo.config.ConfigInstance",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "final"
+    ],
+    "methods" : [
+      "public static java.lang.String getDefMd5()",
+      "public static java.lang.String getDefName()",
+      "public static java.lang.String getDefNamespace()",
+      "public void <init>(ai.vespa.llm.clients.LlmClientConfig$Builder)",
+      "public java.lang.String apiKeySecretName()",
+      "public java.lang.String endpoint()"
+    ],
+    "fields" : [
+      "public static final java.lang.String CONFIG_DEF_MD5",
+      "public static final java.lang.String CONFIG_DEF_NAME",
+      "public static final java.lang.String CONFIG_DEF_NAMESPACE",
+      "public static final java.lang.String[] CONFIG_DEF_SCHEMA"
+    ]
+  },
   "ai.vespa.llm.clients.OpenAI" : {
     "superClass" : "ai.vespa.llm.clients.ConfigurableLanguageModel",
     "interfaces" : [ ],
@@ -9175,30 +9236,102 @@
       "public"
     ],
     "methods" : [
-      "public void <init>(ai.vespa.llm.LlmClientConfig, com.yahoo.container.jdisc.secretstore.SecretStore)",
+      "public void <init>(ai.vespa.llm.clients.LlmClientConfig, com.yahoo.container.jdisc.secretstore.SecretStore)",
       "public java.util.List complete(ai.vespa.llm.completion.Prompt, ai.vespa.llm.InferenceParameters)",
       "public java.util.concurrent.CompletableFuture completeAsync(ai.vespa.llm.completion.Prompt, ai.vespa.llm.InferenceParameters, java.util.function.Consumer)"
     ],
     "fields" : [ ]
   },
-  "ai.vespa.llm.search.LLMSearcher" : {
+  "ai.vespa.search.llm.LLMSearcher" : {
     "superClass" : "com.yahoo.search.Searcher",
     "interfaces" : [ ],
     "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public void <init>(ai.vespa.search.llm.LlmSearcherConfig, com.yahoo.component.provider.ComponentRegistry)",
+      "public com.yahoo.search.Result search(com.yahoo.search.Query, com.yahoo.search.searchchain.Execution)",
+      "protected com.yahoo.search.Result complete(com.yahoo.search.Query, ai.vespa.llm.completion.Prompt)",
+      "public java.lang.String getPrompt(com.yahoo.search.Query)",
+      "public java.lang.String getPropertyPrefix()",
+      "public java.lang.String lookupProperty(java.lang.String, com.yahoo.search.Query)",
+      "public java.lang.Boolean lookupPropertyBool(java.lang.String, com.yahoo.search.Query, boolean)",
+      "public java.lang.String lookupPropertyWithOrWithoutPrefix(java.lang.String, java.util.function.Function)",
+      "public java.lang.String getApiKeyHeader(com.yahoo.search.Query)"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.search.llm.LlmSearcherConfig$Builder" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "com.yahoo.config.ConfigInstance$Builder"
+    ],
+    "attributes" : [
       "public",
+      "final"
+    ],
+    "methods" : [
+      "public void <init>()",
+      "public void <init>(ai.vespa.search.llm.LlmSearcherConfig)",
+      "public ai.vespa.search.llm.LlmSearcherConfig$Builder propertyPrefix(java.lang.String)",
+      "public ai.vespa.search.llm.LlmSearcherConfig$Builder stream(boolean)",
+      "public ai.vespa.search.llm.LlmSearcherConfig$Builder providerId(java.lang.String)",
+      "public final boolean dispatchGetConfig(com.yahoo.config.ConfigInstance$Producer)",
+      "public final java.lang.String getDefMd5()",
+      "public final java.lang.String getDefName()",
+      "public final java.lang.String getDefNamespace()",
+      "public final boolean getApplyOnRestart()",
+      "public final void setApplyOnRestart(boolean)",
+      "public ai.vespa.search.llm.LlmSearcherConfig build()"
+    ],
+    "fields" : [ ]
+  },
+  "ai.vespa.search.llm.LlmSearcherConfig$Producer" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "com.yahoo.config.ConfigInstance$Producer"
+    ],
+    "attributes" : [
+      "public",
+      "interface",
       "abstract"
     ],
-    "methods" : [ ],
+    "methods" : [
+      "public abstract void getConfig(ai.vespa.search.llm.LlmSearcherConfig$Builder)"
+    ],
     "fields" : [ ]
   },
-  "ai.vespa.llm.search.RAGSearcher" : {
-    "superClass" : "ai.vespa.llm.search.LLMSearcher",
+  "ai.vespa.search.llm.LlmSearcherConfig" : {
+    "superClass" : "com.yahoo.config.ConfigInstance",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "final"
+    ],
+    "methods" : [
+      "public static java.lang.String getDefMd5()",
+      "public static java.lang.String getDefName()",
+      "public static java.lang.String getDefNamespace()",
+      "public void <init>(ai.vespa.search.llm.LlmSearcherConfig$Builder)",
+      "public java.lang.String propertyPrefix()",
+      "public boolean stream()",
+      "public java.lang.String providerId()"
+    ],
+    "fields" : [
+      "public static final java.lang.String CONFIG_DEF_MD5",
+      "public static final java.lang.String CONFIG_DEF_NAME",
+      "public static final java.lang.String CONFIG_DEF_NAMESPACE",
+      "public static final java.lang.String[] CONFIG_DEF_SCHEMA"
+    ]
+  },
+  "ai.vespa.search.llm.RAGSearcher" : {
+    "superClass" : "ai.vespa.search.llm.LLMSearcher",
     "interfaces" : [ ],
     "attributes" : [
       "public"
     ],
     "methods" : [
-      "public void <init>(ai.vespa.llm.LlmSearcherConfig, com.yahoo.component.provider.ComponentRegistry)",
+      "public void <init>(ai.vespa.search.llm.LlmSearcherConfig, com.yahoo.component.provider.ComponentRegistry)",
       "public com.yahoo.search.Result search(com.yahoo.search.Query, com.yahoo.search.searchchain.Execution)",
       "protected ai.vespa.llm.completion.Prompt buildPrompt(com.yahoo.search.Query, com.yahoo.search.Result)"
     ],
diff --git a/container-search/src/main/java/ai/vespa/llm/clients/ConfigurableLanguageModel.java b/container-search/src/main/java/ai/vespa/llm/clients/ConfigurableLanguageModel.java
index 662d73d4e01..761fdf0af93 100644
--- a/container-search/src/main/java/ai/vespa/llm/clients/ConfigurableLanguageModel.java
+++ b/container-search/src/main/java/ai/vespa/llm/clients/ConfigurableLanguageModel.java
@@ -3,7 +3,6 @@ package ai.vespa.llm.clients;
 
 import ai.vespa.llm.InferenceParameters;
 import ai.vespa.llm.LanguageModel;
-import ai.vespa.llm.LlmClientConfig;
 import com.yahoo.api.annotations.Beta;
 import com.yahoo.component.annotation.Inject;
 import com.yahoo.container.jdisc.secretstore.SecretStore;
@@ -19,7 +18,7 @@ import java.util.logging.Logger;
 @Beta
 public abstract class ConfigurableLanguageModel implements LanguageModel {
 
-    private static Logger log = Logger.getLogger(ai.vespa.llm.clients.ConfigurableLanguageModel.class.getName());
+    private static Logger log = Logger.getLogger(ConfigurableLanguageModel.class.getName());
 
     private final String apiKey;
     private final String endpoint;
@@ -68,7 +67,9 @@ public abstract class ConfigurableLanguageModel implements LanguageModel {
     }
 
     protected void setEndpoint(InferenceParameters params) {
-        params.setEndpoint(endpoint);
+        if (endpoint != null && ! endpoint.isEmpty()) {
+            params.setEndpoint(endpoint);
+        }
     }
 
 }
diff --git a/container-search/src/main/java/ai/vespa/llm/clients/OpenAI.java b/container-search/src/main/java/ai/vespa/llm/clients/OpenAI.java
index f6092f51948..82e19d47c92 100644
--- a/container-search/src/main/java/ai/vespa/llm/clients/OpenAI.java
+++ b/container-search/src/main/java/ai/vespa/llm/clients/OpenAI.java
@@ -2,7 +2,6 @@
 package ai.vespa.llm.clients;
 
 import ai.vespa.llm.InferenceParameters;
-import ai.vespa.llm.LlmClientConfig;
 import ai.vespa.llm.client.openai.OpenAiClient;
 import ai.vespa.llm.completion.Completion;
 import ai.vespa.llm.completion.Prompt;
diff --git a/container-search/src/main/java/ai/vespa/llm/search/LLMSearcher.java b/container-search/src/main/java/ai/vespa/search/llm/LLMSearcher.java
index 6ff40401a8f..860fc69af91 100755
--- a/container-search/src/main/java/ai/vespa/llm/search/LLMSearcher.java
+++ b/container-search/src/main/java/ai/vespa/search/llm/LLMSearcher.java
@@ -1,12 +1,12 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package ai.vespa.llm.search;
+package ai.vespa.search.llm;
 
 import ai.vespa.llm.InferenceParameters;
 import ai.vespa.llm.LanguageModel;
 import ai.vespa.llm.LanguageModelException;
-import ai.vespa.llm.LlmSearcherConfig;
 import ai.vespa.llm.completion.Completion;
 import ai.vespa.llm.completion.Prompt;
+import ai.vespa.llm.completion.StringPrompt;
 import com.yahoo.api.annotations.Beta;
 import com.yahoo.component.ComponentId;
 import com.yahoo.component.annotation.Inject;
@@ -17,6 +17,7 @@ import com.yahoo.search.Searcher;
 import com.yahoo.search.result.ErrorMessage;
 import com.yahoo.search.result.EventStream;
 import com.yahoo.search.result.HitGroup;
+import com.yahoo.search.searchchain.Execution;
 
 import java.util.List;
 import java.util.function.Function;
@@ -29,7 +30,7 @@ import java.util.stream.Collectors;
  * @author lesters
  */
 @Beta
-public abstract class LLMSearcher extends Searcher {
+public class LLMSearcher extends Searcher {
 
     private static Logger log = Logger.getLogger(LLMSearcher.class.getName());
 
@@ -43,13 +44,18 @@ public abstract class LLMSearcher extends Searcher {
     private final String languageModelId;
 
     @Inject
-    LLMSearcher(LlmSearcherConfig config, ComponentRegistry<LanguageModel> languageModels) {
+    public LLMSearcher(LlmSearcherConfig config, ComponentRegistry<LanguageModel> languageModels) {
         this.stream = config.stream();
         this.languageModelId = config.providerId();
         this.languageModel = findLanguageModel(languageModelId, languageModels);
         this.propertyPrefix = config.propertyPrefix();
     }
 
+    @Override
+    public Result search(Query query, Execution execution) {
+        return complete(query, StringPrompt.from(getPrompt(query)));
+    }
+
     private LanguageModel findLanguageModel(String providerId, ComponentRegistry<LanguageModel> languageModels)
             throws IllegalArgumentException
     {
@@ -74,7 +80,7 @@ public abstract class LLMSearcher extends Searcher {
         return languageModel;
     }
 
-    Result complete(Query query, Prompt prompt) {
+    protected Result complete(Query query, Prompt prompt) {
         var options = new InferenceParameters(getApiKeyHeader(query), s -> lookupProperty(s, query));
         var stream = lookupPropertyBool(STREAM_PROPERTY, query, this.stream);  // query value overwrites config
         return stream ? completeAsync(query, prompt, options) : completeSync(query, prompt, options);
@@ -122,7 +128,7 @@ public abstract class LLMSearcher extends Searcher {
         return new Result(query, hitGroup);
     }
 
-    String getPrompt(Query query) {
+    public String getPrompt(Query query) {
         // Look for prompt with or without prefix
         String prompt = lookupPropertyWithOrWithoutPrefix(PROMPT_PROPERTY, p -> query.properties().getString(p));
         if (prompt != null)
@@ -138,28 +144,28 @@ public abstract class LLMSearcher extends Searcher {
                 "'" + propertyPrefix + "." + PROMPT_PROPERTY + "', '" + PROMPT_PROPERTY + "' or '@query'.");
     }
 
-    String getPropertyPrefix() {
+    public String getPropertyPrefix() {
         return this.propertyPrefix;
     }
 
-    String lookupProperty(String property, Query query) {
+    public String lookupProperty(String property, Query query) {
         String propertyWithPrefix = this.propertyPrefix + "." + property;
         return query.properties().getString(propertyWithPrefix, null);
     }
 
-    Boolean lookupPropertyBool(String property, Query query, boolean defaultValue) {
+    public Boolean lookupPropertyBool(String property, Query query, boolean defaultValue) {
         String propertyWithPrefix = this.propertyPrefix + "." + property;
         return query.properties().getBoolean(propertyWithPrefix, defaultValue);
     }
 
-    String lookupPropertyWithOrWithoutPrefix(String property, Function<String, String> lookup) {
+    public String lookupPropertyWithOrWithoutPrefix(String property, Function<String, String> lookup) {
         String value = lookup.apply(getPropertyPrefix() + "." + property);
         if (value != null)
             return value;
         return lookup.apply(property);
     }
 
-    String getApiKeyHeader(Query query) {
+    public String getApiKeyHeader(Query query) {
         return lookupPropertyWithOrWithoutPrefix(API_KEY_HEADER, p -> query.getHttpRequest().getHeader(p));
     }
 
diff --git a/container-search/src/main/java/ai/vespa/llm/search/RAGSearcher.java b/container-search/src/main/java/ai/vespa/search/llm/RAGSearcher.java
index b8e33778ced..cba153d881d 100755
--- a/container-search/src/main/java/ai/vespa/llm/search/RAGSearcher.java
+++ b/container-search/src/main/java/ai/vespa/search/llm/RAGSearcher.java
@@ -1,8 +1,7 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package ai.vespa.llm.search;
+package ai.vespa.search.llm;
 
 import ai.vespa.llm.LanguageModel;
-import ai.vespa.llm.LlmSearcherConfig;
 import ai.vespa.llm.completion.Prompt;
 import ai.vespa.llm.completion.StringPrompt;
 import com.yahoo.api.annotations.Beta;
diff --git a/container-search/src/main/java/ai/vespa/llm/search/package-info.java b/container-search/src/main/java/ai/vespa/search/llm/package-info.java
index 6a8975fd2fa..60af24c47b2 100644
--- a/container-search/src/main/java/ai/vespa/llm/search/package-info.java
+++ b/container-search/src/main/java/ai/vespa/search/llm/package-info.java
@@ -1,7 +1,7 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 @ExportPackage
 @PublicApi
-package ai.vespa.llm.search;
+package ai.vespa.search.llm;
 
 import com.yahoo.api.annotations.PublicApi;
 import com.yahoo.osgi.annotation.ExportPackage;
diff --git a/container-search/src/main/resources/configdefinitions/llm-client.def b/container-search/src/main/resources/configdefinitions/llm-client.def
index 6bfd95c3cf2..0866459166a 100755
--- a/container-search/src/main/resources/configdefinitions/llm-client.def
+++ b/container-search/src/main/resources/configdefinitions/llm-client.def
@@ -1,5 +1,5 @@
 # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package=ai.vespa.llm
+package=ai.vespa.llm.clients
 
 # The name of the secret containing the api key
 apiKeySecretName string default=""
diff --git a/container-search/src/main/resources/configdefinitions/llm-searcher.def b/container-search/src/main/resources/configdefinitions/llm-searcher.def
index 918a6e6e8b1..608f44f04d0 100755
--- a/container-search/src/main/resources/configdefinitions/llm-searcher.def
+++ b/container-search/src/main/resources/configdefinitions/llm-searcher.def
@@ -1,5 +1,5 @@
 # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package=ai.vespa.llm
+package=ai.vespa.search.llm
 
 # Query propertry prefix for options
 propertyPrefix string default="llm"
diff --git a/container-search/src/test/java/ai/vespa/llm/clients/ConfigurableLanguageModelTest.java b/container-search/src/test/java/ai/vespa/llm/clients/ConfigurableLanguageModelTest.java
index 1f2a12322a1..a9f4c3dfac5 100644
--- a/container-search/src/test/java/ai/vespa/llm/clients/ConfigurableLanguageModelTest.java
+++ b/container-search/src/test/java/ai/vespa/llm/clients/ConfigurableLanguageModelTest.java
@@ -2,7 +2,6 @@
 package ai.vespa.llm.clients;
 
 import ai.vespa.llm.InferenceParameters;
-import ai.vespa.llm.LlmClientConfig;
 import ai.vespa.llm.completion.Completion;
 import ai.vespa.llm.completion.Prompt;
 import ai.vespa.llm.completion.StringPrompt;
diff --git a/container-search/src/test/java/ai/vespa/llm/clients/MockLLMClient.java b/container-search/src/test/java/ai/vespa/llm/clients/MockLLMClient.java
index cfb6a43984f..4d0073f1cbe 100644
--- a/container-search/src/test/java/ai/vespa/llm/clients/MockLLMClient.java
+++ b/container-search/src/test/java/ai/vespa/llm/clients/MockLLMClient.java
@@ -2,7 +2,6 @@
 package ai.vespa.llm.clients;
 
 import ai.vespa.llm.InferenceParameters;
-import ai.vespa.llm.LlmClientConfig;
 import ai.vespa.llm.completion.Completion;
 import ai.vespa.llm.completion.Prompt;
 import com.yahoo.container.jdisc.secretstore.SecretStore;
diff --git a/container-search/src/test/java/ai/vespa/llm/clients/OpenAITest.java b/container-search/src/test/java/ai/vespa/llm/clients/OpenAITest.java
index 1111a9824f5..57339f6ad49 100644
--- a/container-search/src/test/java/ai/vespa/llm/clients/OpenAITest.java
+++ b/container-search/src/test/java/ai/vespa/llm/clients/OpenAITest.java
@@ -2,7 +2,6 @@
 package ai.vespa.llm.clients;
 
 import ai.vespa.llm.InferenceParameters;
-import ai.vespa.llm.LlmClientConfig;
 import ai.vespa.llm.completion.StringPrompt;
 import com.yahoo.container.jdisc.SecretStoreProvider;
 import org.junit.jupiter.api.Disabled;
diff --git a/container-search/src/test/java/ai/vespa/llm/search/LLMSearcherTest.java b/container-search/src/test/java/ai/vespa/search/llm/LLMSearcherTest.java
index acdfaf9fffc..1efcf1c736a 100755
--- a/container-search/src/test/java/ai/vespa/llm/search/LLMSearcherTest.java
+++ b/container-search/src/test/java/ai/vespa/search/llm/LLMSearcherTest.java
@@ -1,14 +1,12 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package ai.vespa.llm.search;
+package ai.vespa.search.llm;
 
 import ai.vespa.llm.InferenceParameters;
 import ai.vespa.llm.LanguageModel;
-import ai.vespa.llm.LlmClientConfig;
-import ai.vespa.llm.LlmSearcherConfig;
 import ai.vespa.llm.clients.ConfigurableLanguageModelTest;
+import ai.vespa.llm.clients.LlmClientConfig;
 import ai.vespa.llm.clients.MockLLMClient;
 import ai.vespa.llm.completion.Prompt;
-import ai.vespa.llm.completion.StringPrompt;
 import com.yahoo.component.ComponentId;
 import com.yahoo.component.chain.Chain;
 import com.yahoo.component.provider.ComponentRegistry;
@@ -229,26 +227,14 @@ public class LLMSearcherTest {
         ComponentRegistry<LanguageModel> models = new ComponentRegistry<>();
         llms.forEach((key, value) -> models.register(ComponentId.fromString(key), value));
         models.freeze();
-        return new LLMSearcherImpl(config, models);
+        return new LLMSearcher(config, models);
     }
 
     private static Searcher createLLMSearcher(LlmSearcherConfig config, Map<String, LanguageModel> llms) {
         ComponentRegistry<LanguageModel> models = new ComponentRegistry<>();
         llms.forEach((key, value) -> models.register(ComponentId.fromString(key), value));
         models.freeze();
-        return new LLMSearcherImpl(config, models);
-    }
-
-    public static class LLMSearcherImpl extends LLMSearcher {
-
-        public LLMSearcherImpl(LlmSearcherConfig config, ComponentRegistry<LanguageModel> languageModels) {
-            super(config, languageModels);
-        }
-
-        @Override
-        public Result search(Query query, Execution execution) {
-            return complete(query, StringPrompt.from(getPrompt(query)));
-        }
+        return new LLMSearcher(config, models);
     }
 
 }
diff --git a/container-search/src/test/java/ai/vespa/llm/search/RAGSearcherTest.java b/container-search/src/test/java/ai/vespa/search/llm/RAGSearcherTest.java
index ccf9a4a6401..d6b66b1a8c6 100755
--- a/container-search/src/test/java/ai/vespa/llm/search/RAGSearcherTest.java
+++ b/container-search/src/test/java/ai/vespa/search/llm/RAGSearcherTest.java
@@ -1,8 +1,7 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package ai.vespa.llm.search;
+package ai.vespa.search.llm;
 
 import ai.vespa.llm.LanguageModel;
-import ai.vespa.llm.LlmSearcherConfig;
 import com.yahoo.component.ComponentId;
 import com.yahoo.component.chain.Chain;
 import com.yahoo.component.provider.ComponentRegistry;
diff --git a/vespajlib/src/test/java/ai/vespa/llm/client/openai/OpenAiClientCompletionTest.java b/vespajlib/src/test/java/ai/vespa/llm/client/openai/OpenAiClientCompletionTest.java
index 1baab26f496..4dd02f01d05 100644
--- a/vespajlib/src/test/java/ai/vespa/llm/client/openai/OpenAiClientCompletionTest.java
+++ b/vespajlib/src/test/java/ai/vespa/llm/client/openai/OpenAiClientCompletionTest.java
@@ -2,7 +2,6 @@
 package ai.vespa.llm.client.openai;
 
 import ai.vespa.llm.InferenceParameters;
-import ai.vespa.llm.completion.Completion;
 import ai.vespa.llm.completion.StringPrompt;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
diff --git a/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java b/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java
index 24c496a3d2c..d31a5129156 100644
--- a/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java
+++ b/vespajlib/src/test/java/ai/vespa/llm/completion/CompletionTest.java
@@ -1,8 +1,8 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.llm.completion;
 
-import ai.vespa.llm.InferenceParameters;
 import ai.vespa.llm.test.MockLanguageModel;
+import ai.vespa.llm.InferenceParameters;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;