Glebashnik/feed field generator

indexinglanguage/src/main/java/com/yahoo/vespa/indexinglanguage/expressions/GenerateExpression.java

@@ -82,7 +82,7 @@ public DataType setOutputType(DataType outputType, VerificationContext context)
             throw new VerificationException(this, "Generate expression requires either a string or array<string> output type, but got "
                     + outputType.getName());
 
-        super.setOutputType(null, outputType, null, context); // todo: Why not set actualOutput to outputType?
+        super.setOutputType(null, outputType, null, context); // TODO: Why not set actualOutput to outputType?
         return outputType; // return input type the same as output type: string or array<string>
     }
 

model-integration/src/main/java/ai/vespa/llm/clients/LocalLLM.java

@@ -8,7 +8,6 @@
 import ai.vespa.llm.completion.Prompt;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.component.annotation.Inject;
-import com.yahoo.language.process.TextGenerator;
 import de.kherud.llama.LlamaModel;
 import de.kherud.llama.ModelParameters;
 

model-integration/src/main/java/ai/vespa/llm/generation/LanguageModelTextGenerator.java

@@ -3,11 +3,15 @@
 import ai.vespa.llm.InferenceParameters;
 import ai.vespa.llm.LanguageModel;
 import ai.vespa.llm.completion.Prompt;
+import ai.vespa.llm.completion.StringPrompt;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.component.annotation.Inject;
 import com.yahoo.component.provider.ComponentRegistry;
 import com.yahoo.language.process.TextGenerator;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.logging.Logger;
 
 /**
@@ -18,17 +22,50 @@
 public class LanguageModelTextGenerator extends AbstractComponent implements TextGenerator {
     private static final Logger logger = Logger.getLogger(LanguageModelTextGenerator.class.getName());
     private final LanguageModel languageModel;
+
+    // Template usually contains {input} placeholder for the dynamic part of the prompt, replaced with the actual input.
+    // Templates without {input} are possible, which will ignore the input, making the prompt static.
+    // TODO: Consider not allowing templates without {input} to avoid costly errors when users forget to include {input}.
+    private static final String DEFAULT_PROMPT_TEMPLATE = "{input}";
+    private final String promptTemplate;
 
     @Inject
-    public LanguageModelTextGenerator(TextGeneratorConfig config, ComponentRegistry<LanguageModel> languageModels) {
+    public LanguageModelTextGenerator(LanguageModelTextGeneratorConfig config, ComponentRegistry<LanguageModel> languageModels) {
         this.languageModel = LanguageModelUtils.findLanguageModel(config.providerId(), languageModels, logger);
+        this.promptTemplate = loadPromptTemplate(config);
     }
 
+    private String loadPromptTemplate(LanguageModelTextGeneratorConfig config) {
+        if (config.promptTemplate() != null && !config.promptTemplate().isEmpty()) {
+            return config.promptTemplate();
+        } else if (config.promptTemplateFile().isPresent()) {
+            Path path = config.promptTemplateFile().get();
+
+            try {
+                String promptTemplate = new String(Files.readAllBytes(path));
+
+                if (!promptTemplate.isEmpty()) {  // TODO: Consider throwing an exception if the template is empty.
+                    return promptTemplate;
+                }
+            } catch (IOException e) {
+                throw new IllegalArgumentException("Could not read prompt template file: " + path, e);
+            }
+        }
+
+        return DEFAULT_PROMPT_TEMPLATE;
+    }
+
     @Override
     public String generate(Prompt prompt, Context context) {
+        var finalPrompt = buildPrompt(prompt);
         var options = new InferenceParameters(s -> "");
-        var completions = languageModel.complete(prompt, options);
+        var completions = languageModel.complete(finalPrompt, options);
         var firstCompletion = completions.get(0);
         return firstCompletion.text();
     }
+
+    private Prompt buildPrompt(Prompt inputPrompt) {
+        String finalPrompt = promptTemplate.replace("{input}", inputPrompt.asString());
+        return StringPrompt.from(finalPrompt);
+    }
 }

model-integration/src/main/resources/configdefinitions/language-model-text-generator.def

@@ -1,11 +1,11 @@
 # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package=ai.vespa.llm.generation
 
-# The external LLM provider - the id of a LanguageModel component
-providerId string default=""
+# Id of a LanguageModel component specified in services.xml, e.g. OpenAI, LocalLLM
+providerId string
 
-# The default prompt to use if not overridden in query
-prompt string default=""
+# Prompt template, e.g. "Extract named entities from this text, output one entity per line: {input}"
+promptTemplate string default=""
 
-# The default prompt template file to use if not overridden in query. Above prompt has precedence if it is set.
-promptTemplate path optional
+# Prompt template in a file, above promptTemplate has precedence if it is set.
+promptTemplateFile path optional

model-integration/src/test/java/ai/vespa/llm/generation/LanguageModelTextGeneratorTest.java

@@ -10,35 +10,174 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 
+import com.yahoo.config.FileReference;
 import com.yahoo.language.process.TextGenerator;
 import org.junit.Test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 
 public class LanguageModelTextGeneratorTest {
+
     @Test
     public void testGenerate() {
-        LanguageModel languageModel1 = new RepeaterMockLanguageModel(1);
-        LanguageModel languageModel2 = new RepeaterMockLanguageModel(2);
-        var languageModels = Map.of("mock1", languageModel1, "mock2", languageModel2);
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder().providerId("languageModel").build();
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+        var result = generator.generate(StringPrompt.from("hello"), context);
+        assertEquals("hello hello", result);
+    }
+
+    @Test
+    public void testGenerateWithTwoLanguageModel() {
+        LanguageModel languageModel1 = new RepeaterMockLanguageModel(2);
+        LanguageModel languageModel2 = new RepeaterMockLanguageModel(3);
+        var languageModels = Map.of("languageModel1", languageModel1, "languageModel2", languageModel2);
 
-        var config1 = new TextGeneratorConfig.Builder().providerId("mock1").build();
+        var config1 = new LanguageModelTextGeneratorConfig.Builder().providerId("languageModel1").build();
         var generator1 = createGenerator(config1, languageModels);
         var context = new TextGenerator.Context("schema.indexing");
         var result1 = generator1.generate(StringPrompt.from("hello"), context);
-        assertEquals("hello", result1);
+        assertEquals("hello hello", result1);
 
-        var config2 = new TextGeneratorConfig.Builder().providerId("mock2").build();
-        var generator2 = createGenerator(config2, Map.of("mock1", languageModel1, "mock2", languageModel2));
+        var config2 = new LanguageModelTextGeneratorConfig.Builder().providerId("languageModel2").build();
+        var generator2 = createGenerator(config2, languageModels);
         var result2 = generator2.generate(StringPrompt.from("hello"), context);
+        assertEquals("hello hello hello", result2);
+    }
+
+    @Test
+    public void testGenerateWithPromptTemplate() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplate("hello {input}")
+                .build();
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("hello world hello world", result1);
+
+        var result2 = generator.generate(StringPrompt.from("there"), context);
+        assertEquals("hello there hello there", result2);
+    }
+
+    @Test
+    public void testGenerateWithEmptyPromptTemplate() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplate("")
+                .build();
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("world world", result1);
+    }
+
+    @Test
+    public void testGenerateWithStaticPromptTemplate() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplate("hello")
+                .build();
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("hello hello", result1);
+
+        var result2 = generator.generate(StringPrompt.from("there"), context);
         assertEquals("hello hello", result2);
     }
+
+    @Test
+    public void testGenerateWithPromptTemplateFile() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplateFile(Optional.of(new FileReference("src/test/prompts/prompt_with_input.txt")))
+                .build();
+
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("hello world hello world", result1);
+    }
+
+    @Test
+    public void testGenerateWithEmptyTemplateFile() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplateFile(Optional.of(new FileReference("src/test/prompts/empty_prompt.txt")))
+                .build();
+
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("world world", result1);
+    }
+
+    @Test
+    public void testGenerateWithMissingTemplateFile() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplateFile(Optional.of(new FileReference("src/test/prompts/missing_prompt.txt")))
+                .build();
+
+        try {
+            createGenerator(config, languageModels);
+        }
+        catch (IllegalArgumentException e) {
+            assertEquals("Could not read prompt template file: src/test/prompts/missing_prompt.txt", e.getMessage());
+        }
+    }
+
+    @Test
+    public void testGenerateWithPromptTemplateOverridesPromptTemplateFile() {
+        LanguageModel languageModel = new RepeaterMockLanguageModel(2);
+        var languageModels = Map.of("languageModel", languageModel);
+
+        var config = new LanguageModelTextGeneratorConfig.Builder()
+                .providerId("languageModel")
+                .promptTemplate("bye {input}")
+                .promptTemplateFile(Optional.of(new FileReference("src/test/prompts/prompt_with_input.txt")))
+                .build();
+
+        var generator = createGenerator(config, languageModels);
+        var context = new TextGenerator.Context("schema.indexing");
+
+        var result1 = generator.generate(StringPrompt.from("world"), context);
+        assertEquals("bye world bye world", result1);
+    }
 
-    private static LanguageModelTextGenerator createGenerator(TextGeneratorConfig config, Map<String, LanguageModel> languageModels) {
+    private static LanguageModelTextGenerator createGenerator(LanguageModelTextGeneratorConfig config, Map<String, LanguageModel> languageModels) {
         ComponentRegistry<LanguageModel> languageModelsRegistry = new ComponentRegistry<>();
         languageModels.forEach((key, value) -> languageModelsRegistry.register(ComponentId.fromString(key), value));
         languageModelsRegistry.freeze();

model-integration/src/test/prompts/prompt_with_input.txt

@@ -0,0 +1 @@
+hello {input}