diff options
author | Jon Bratseth <bratseth@gmail.com> | 2023-04-26 14:38:13 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-04-26 14:38:13 +0200 |
commit | 072acd73475aae1a882d908d57d28eec31e69c3e (patch) | |
tree | 33ec28c28ff9e2d0badf68df94fff9ec05df5c04 /vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java | |
parent | 5b290735b485e20216fac0733c0d44274eec3937 (diff) | |
parent | 901f992bf024fe1b994385ca76b7b59fe2d8a78f (diff) |
Merge pull request #26851 from vespa-engine/bratseth/openai-component
Put the openai client in a separate component
Diffstat (limited to 'vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java')
-rw-r--r-- | vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java b/vespajlib/src/main/java/ai/vespa/llm/test/MockLanguageModel.java new file mode 100644 index 00000000000..d47f43c55b2 --- /dev/null +++ b/vespajlib/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.LanguageModel; +import ai.vespa.llm.completion.Completion; +import ai.vespa.llm.completion.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); } + + } + +} |