diff options
author | Jon Bratseth <bratseth@vespa.ai> | 2023-04-19 11:21:47 +0200 |
---|---|---|
committer | Jon Bratseth <bratseth@vespa.ai> | 2023-04-19 11:21:47 +0200 |
commit | eaccd0e3918016013f426eecdb7e0fc3493cece9 (patch) | |
tree | 0e209fb66c44bb29a9125e1a3f16cf3d16eceea6 /model-integration/src/main | |
parent | 2d9fd2e6c78e9ae8580cd4a21d20d0febb8f9c93 (diff) |
Avoid ignored directory name
Diffstat (limited to 'model-integration/src/main')
-rw-r--r-- | model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java b/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java new file mode 100644 index 00000000000..54b085a451c --- /dev/null +++ b/model-integration/src/main/java/ai/vespa/llm/test/MockLanguageModel.java @@ -0,0 +1,44 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +package ai.vespa.llm.test; + +import ai.vespa.llm.Completion; +import ai.vespa.llm.LanguageModel; +import ai.vespa.llm.Prompt; +import com.yahoo.api.annotations.Beta; + +import java.util.List; +import java.util.function.Function; + +/** + * @author bratseth + */ +@Beta +public class MockLanguageModel implements LanguageModel { + + private final Function<Prompt, List<Completion>> completer; + + public MockLanguageModel(Builder builder) { + completer = builder.completer; + } + + @Override + public List<Completion> complete(Prompt prompt) { + return completer.apply(prompt); + } + + public static class Builder { + + private Function<Prompt, List<Completion>> completer = prompt -> List.of(Completion.from("")); + + public Builder completer(Function<Prompt, List<Completion>> completer) { + this.completer = completer; + return this; + } + + public Builder() {} + + public MockLanguageModel build() { return new MockLanguageModel(this); } + + } + +} |